Update modules/translation/translation_base.py

modules/translation/translation_base.py

@@ -185,7 +185,16 @@ class TranslationBase(ABC):
             raise
         finally:
             self.release_cuda_memory()
-    
+
+    def offload(self):
+        """Offload the model and free up the memory"""
+        if self.model is not None:
+            del self.model
+            self.model = None
+        if self.device == "cuda":
+            self.release_cuda_memory()
+        gc.collect()
+
     @staticmethod
     def get_device():
         if torch.cuda.is_available():