Fix bug in visualization callback
This commit is contained in:
parent
3148684812
commit
8d57f69c9e
@ -27,6 +27,7 @@ class VisualizationCallback(pl.Callback):
|
||||
def on_epoch_end(self, trainer, pl_module):
|
||||
protos = pl_module.prototypes
|
||||
plabels = pl_module.prototype_labels
|
||||
x_train, y_train = self.x_train, self.y_train
|
||||
ax = self.fig.gca()
|
||||
ax.cla()
|
||||
ax.set_title(self.title)
|
||||
|
Loading…
Reference in New Issue
Block a user