fix: handling cases where pytorch file is a single tensor

dreadnode · Nov 8, 2024 · 14d2d79 · 14d2d79
1 parent c40b642
commit 14d2d79
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/core/handlers/pytorch/inspect.py b/src/core/handlers/pytorch/inspect.py
@@ -32,6 +32,10 @@ def main():
     all_metadata = getattr(model, "_metadata") if hasattr(model, "_metadata") else {}
     model_metadata = all_metadata[""] if "" in all_metadata else {}
 
+    # some files are just a single tensor
+    if isinstance(model, torch.Tensor):
+        model = {"<unamed-single-tensor>": model}
+
     inspection = {
         "file_path": file_path,
         "file_type": "PyTorch",