Add kwargs to weight descr of bioimageio

ilastik · Jan 20, 2025 · 6e4ec16 · 6e4ec16
1 parent 82b4aba
commit 6e4ec16
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tiktorch/trainer.py b/tiktorch/trainer.py
@@ -310,6 +310,7 @@ def export(self, file_to_save: Path):
                     architecture=ArchitectureFromLibraryDescr(
                         import_from=f"{self.get_model_import_file_path()}",
                         callable=Identifier(f"{self.model.__class__.__name__}"),
+                        kwargs={"in_channels": self._in_channels, "out_channels": self._out_channels},
                     ),
                     pytorch_version=Version("1.1.1"),
                 )