Skip to content

Commit 01938fb

Browse files
committed
Deploy memory usage fix in inference as well
1 parent 7bb5edc commit 01938fb

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

napari_cellseg3d/code_models/worker_inference.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -784,6 +784,12 @@ def inference(self):
784784
)
785785

786786
model.to("cpu")
787+
model = None
788+
del model
789+
inference_loader = None
790+
del inference_loader
791+
if torch.cuda.is_available():
792+
torch.cuda.empty_cache()
787793
# self.quit()
788794
except Exception as e:
789795
logger.exception(e)

0 commit comments

Comments
 (0)