Improvements

onnx · Mar 1, 2024 · 40e0efd · 40e0efd
1 parent 34ea505
commit 40e0efd
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/turnkeyml/analyze/status.py b/src/turnkeyml/analyze/status.py
@@ -335,6 +335,9 @@ def add_unique_invocation(
         parent_hash: Union[str, None] = None,
         executed: int = 0,
     ):
+        model_class = (
+            type(self.model) if self.model_type == build.ModelType.PYTORCH else None
+        )
         self.unique_invocations[invocation_hash] = UniqueInvocationInfo(
             name=self.name,
             script_name=self.script_name,
@@ -344,7 +347,7 @@ def add_unique_invocation(
             depth=self.depth,
             build_model=self.build_model,
             model_type=self.model_type,
-            model_class=type(self.model),
+            model_class=model_class,
             invocation_hash=invocation_hash,
             hash=self.hash,
             is_target=is_target,