From a3e183d7fc5946467da3bb6430c565fd03967dfb Mon Sep 17 00:00:00 2001 From: andreapdr Date: Tue, 14 Mar 2023 11:22:00 +0100 Subject: [PATCH] avoid duplicating model on gpu when earlystop is triggered --- gfun/vgfs/commons.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/gfun/vgfs/commons.py b/gfun/vgfs/commons.py index b1f01cd..e6527cd 100644 --- a/gfun/vgfs/commons.py +++ b/gfun/vgfs/commons.py @@ -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: