avoid duplicating model on gpu when earlystop is triggered

This commit is contained in:
Andrea Pedrotti 2023-03-14 11:22:00 +01:00
parent 57918ec523
commit a3e183d7fc
1 changed files with 6 additions and 3 deletions

View File

@ -243,9 +243,12 @@ class Trainer:
print(
f"- restoring best model from epoch {self.earlystopping.best_epoch} with best metric: {self.earlystopping.best_score:3f}"
)
self.model = self.earlystopping.load_model(self.model).to(
self.device
)
restored_model = self.earlystopping.load_model(self.model)
# swapping model on gpu
del self.model
self.model = restored_model.to(self.device)
break
if self.scheduler is not None: