fix bug

EleutherAI · Jan 27, 2024 · 2233495 · 2233495
1 parent 830daa9
commit 2233495
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/aria/utils.py b/aria/utils.py
@@ -71,4 +71,6 @@ 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)