Removed helper method

nygaardcodecommentclassification/api/models.py

@@ -151,19 +151,9 @@ class ModelRegistry:
                 lang.upper(),
                 catboost_run_name,
             )
-            # Configure torch to map GPU tensors to CPU during deserialization
-            original_map_location = getattr(torch, '_utils', {}).get('_rebuild_device_tensor_from_numpy', None)
-            try:
-                # Set the map_location handler for torch.load to handle GPU-trained models
-                model = mlflow.sklearn.load_model(model_uri)
-            except Exception as load_error:
-                logger.warning(
-                    "[%s] Initial model load failed, attempting CPU mapping: %s",
-                    lang.upper(),
-                    load_error
-                )
-                # Fallback: attempt to patch torch to use CPU for deserialization
-                model = self._load_model_with_cpu_mapping(model_uri, lang)
+            # Set default torch device to CPU before loading
+            torch.set_default_device("cpu")
+            model = mlflow.sklearn.load_model(model_uri)
 
             # Load the sentence transformer embedder from MLflow
             embedder_uri = f"runs:/{embedder_run_id}/model_{lang}"
@@ -172,16 +162,7 @@ class ModelRegistry:
                 lang.upper(),
                 embedder_run_name,
             )
-            try:
-                embedder = mlflow.sklearn.load_model(embedder_uri)
-            except Exception as embed_error:
-                logger.warning(
-                    "[%s] Initial embedder load failed, attempting CPU mapping: %s",
-                    lang.upper(),
-                    embed_error
-                )
-                embedder = self._load_model_with_cpu_mapping(embedder_uri, lang)
-            
+            embedder = mlflow.sklearn.load_model(embedder_uri)
             if hasattr(embedder, "to"):
                 embedder.to("cpu")
 
@@ -200,32 +181,6 @@ class ModelRegistry:
         except Exception as e:
             logger.error("[%s] Error loading models: %s", lang.upper(), e)
 
-    def _load_model_with_cpu_mapping(self, model_uri: str, lang: str) -> Any:
-        """Load a model with explicit CPU device mapping to handle GPU-trained models.
-
-        This method attempts to load models trained on GPU on CPU-only machines
-        by setting torch's device context and using torch.load with map_location.
-
-        Args:
-            model_uri: MLflow model URI
-            lang: Programming language code for logging
-
-        Returns:
-            Loaded model object mapped to CPU
-
-        Raises:
-            Exception: If all loading attempts fail
-        """
-        # Set torch to use CPU for any tensor operations
-        torch.set_default_device("cpu")
-        
-        try:
-            model = mlflow.sklearn.load_model(model_uri)
-            return model
-        finally:
-            # Reset to default behavior
-            torch.set_default_device(None)
-
     def get_model(self, language: str, model_type: str) -> Optional[Dict[str, Any]]:
         """Retrieve a loaded model entry by language and type.