avoid duplicating model on gpu when earlystop is triggered
This commit is contained in:
parent
57918ec523
commit
a3e183d7fc
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue