From f4bac252ae4ed3555f2ac22744bfb2ec9f6be9e2 Mon Sep 17 00:00:00 2001 From: Michael Kofler Date: Thu, 25 Jun 2026 14:29:46 +0200 Subject: [PATCH 1/3] Refactor model loading to handle module prefix in state_dict more efficiently --- DeepSDFStruct/deep_sdf/workspace.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/DeepSDFStruct/deep_sdf/workspace.py b/DeepSDFStruct/deep_sdf/workspace.py index bb034b39..88ecd5c8 100644 --- a/DeepSDFStruct/deep_sdf/workspace.py +++ b/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 From c893e5645a75dcfbc2dff79132e8fd9402fe5a6a Mon Sep 17 00:00:00 2001 From: Michael Kofler Date: Thu, 25 Jun 2026 14:30:53 +0200 Subject: [PATCH 2/3] added some additional output checks for the model export --- DeepSDFStruct/deep_sdf/models.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/DeepSDFStruct/deep_sdf/models.py b/DeepSDFStruct/deep_sdf/models.py index 4635dc6b..9f87cbdb 100644 --- a/DeepSDFStruct/deep_sdf/models.py +++ b/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}") From 3d4160d37a115fb1b1abe8a1fd67be558e45b444 Mon Sep 17 00:00:00 2001 From: Michael Kofler Date: Thu, 25 Jun 2026 16:03:28 +0200 Subject: [PATCH 3/3] replaced deprecated torch.linalg.norm argument axes --- DeepSDFStruct/deep_sdf/networks/analytic_round_cross.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DeepSDFStruct/deep_sdf/networks/analytic_round_cross.py b/DeepSDFStruct/deep_sdf/networks/analytic_round_cross.py index af959ddc..2f0b8db7 100644 --- a/DeepSDFStruct/deep_sdf/networks/analytic_round_cross.py +++ b/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