Adds option to set model checkpoint path in ETM model
@@ -70,12 +70,12 @@ class ETM(BaseETM):
self.optimizer = None
self.embeddings = None
- def train_model(self, dataset, hyperparameters=None, top_words=10):
+ def train_model(self, dataset, hyperparameters=None, top_words=10, op_path='checkpoint.pt'):
if hyperparameters is None:
hyperparameters = {}
self.set_model(dataset, hyperparameters)
self.top_word = top_words
- self.early_stopping = EarlyStopping(patience=5, verbose=True)
+ self.early_stopping = EarlyStopping(patience=5, verbose=True, path=op_path)
for epoch in range(0, self.hyperparameters['num_epochs']):
continue_training = self._train_epoch(epoch)
