fix bug

EleutherAI · Jan 27, 2024 · e57225d · e57225d
1 parent 830daa9
commit e57225d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/aria/utils.py b/aria/utils.py
@@ -71,4 +71,5 @@ def _load_weight(ckpt_path: str, device="cpu"):
             ) from e
         return load_file(ckpt_path, device=device)
     else:
+        import torch
         return torch.load(ckpt_path, map_location=device)