Fix wrong module loading

DeepSDFStruct/deep_sdf/models.py

@@ -155,12 +155,14 @@ def _decode_sdf(
         model_input = torch.cat([latent_repeat, queries], dim=1)
         return self._decoder(model_input)
 
-    def export_libtorch_executable(self, filename: str):
+    def export_libtorch_executable(self, filename: str, use_script: bool = True):
         """
         Export the trained decoder model to a TorchScript file for use with LibTorch (C++).
 
         Args:
             filename (str): Path where the TorchScript model will be saved (e.g. "decoder.pt").
+            use_script (bool): If True, use torch.jit.script (recommended for C++ compatibility).
+                              If False, use torch.jit.trace (legacy mode, may have control flow issues).
 
         Example:
             >>> model.export_libtorch_executable("decoder.pt")
@@ -176,13 +178,20 @@ def export_libtorch_executable(self, filename: str):
         ), "trained_latent_vectors must contain at least one element"
         latent = self._trained_latent_vectors
         example_input = torch.cat(
-            [latent[0], torch.tensor([0, 0, 0], device=self.device)]
+            [latent[0], torch.tensor([0.0, 0.0, 0.0], device=self.device)]
         ).unsqueeze(0)
 
         print("Example input: ", example_input)
         print("Example Output: ", self._decoder(example_input))
 
-        decoder_traced = torch.jit.trace(self._decoder, example_input)
-        sm = torch.jit.script(decoder_traced)
+        self._decoder.eval()
+
+        if use_script:
+            print("Exporting with torch.jit.script (recommended)")
+            sm = torch.jit.script(self._decoder)
+        else:
+            print("Exporting with torch.jit.trace (legacy mode)")
+            sm = torch.jit.trace(self._decoder, example_input)
 
         sm.save(filename)
+        print(f"Saved to {filename}")

DeepSDFStruct/deep_sdf/networks/analytic_round_cross.py

@@ -90,7 +90,7 @@ def make_sequence():
     def forward(self, input):
         xyz = input[:, -self.geom_dimension :]
         r = input[:, 0]
-        output = torch.linalg.norm(xyz, axis=1, ord=torch.inf)
+        output = torch.linalg.norm(xyz, dim=1, ord=torch.inf)
 
         # add x cylinder
         cylinder = torch.sqrt(xyz[:, 1] ** 2 + xyz[:, 2] ** 2) - r

DeepSDFStruct/deep_sdf/workspace.py

@@ -316,14 +316,16 @@ def load_trained_model(
 
     data = torch.load(filename, map_location=device)
     decoder = init_decoder(experiment_specs, device, data_parallel)
-    try:
-        decoder.load_state_dict(data["model_state_dict"], strict=False)
-    except RuntimeError:
-        state_dict = {}
-        for k, v in data["model_state_dict"].items():
-            new_key = k.replace("module.", "", 1) if k.startswith("module.") else k
-            state_dict[new_key] = v
-        decoder.load_state_dict(state_dict, strict=False)
+
+    state_dict = data["model_state_dict"]
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        stripped_dict = {}
+        for k, v in state_dict.items():
+            new_key = k.replace("module.", "", 1)
+            stripped_dict[new_key] = v
+        state_dict = stripped_dict
+
+    decoder.load_state_dict(state_dict, strict=False)
     decoder = decoder.to(device)
     return decoder