Update src/modules/feature_extractor.py

src/modules/feature_extractor.py

@@ -46,7 +46,7 @@ class FeatureExtractor:
             self.model, _ = self.init_model(base_model)
             self.model.eval()
             self.device = torch.device("cpu")
-            # self.model.to(self.device).half()
+            self.model.to(self.device)
 
             # Create directory if it doesn't exist
             os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
@@ -98,8 +98,7 @@ class FeatureExtractor:
         # add batch dimension
         x = x.unsqueeze(0)
 
-        # Convert to float16 and then to numpy for ONNX Runtime
-        x = x.half()
+        # Convert to numpy for ONNX Runtime
         x_numpy = x.numpy()
         # Get model size in MB
         model_size = os.path.getsize(self.onnx_path) / (1024 * 1024)
@@ -126,10 +125,6 @@ class FeatureExtractor:
         dummy_input = torch.randn(1, 3, 224, 224, device=self.device)
 
         # Export the model
-        # Convert model to float16
-        self.model.half()
-        dummy_input = dummy_input.half()
-        
         torch.onnx.export(
             self.model,
             dummy_input,