+if train_input.size(1) == 1:
+
+ ax.set_xlim(-1.25, 1.25)
+
+ d = train_input.flatten().detach().numpy()
+ ax.hist(d, 25, (-1, 1),
+ density = True,
+ histtype = 'stepfilled', color = 'lightblue', label = 'Train')
+
+ d = x.flatten().detach().numpy()
+ ax.hist(d, 25, (-1, 1),
+ density = True,
+ histtype = 'step', color = 'red', label = 'Synthesis')
+
+ ax.legend(frameon = False, loc = 2)
+
+elif train_input.size(1) == 2:
+
+ ax.set_xlim(-1.25, 1.25)
+ ax.set_ylim(-1.25, 1.25)
+ ax.set(aspect = 1)
+
+ d = train_input[:200].detach().numpy()
+ ax.scatter(d[:, 0], d[:, 1],
+ color = 'lightblue', label = 'Train')