Back to meta

Author: jackzhxng

examples/models/llama/model.py

@@ -171,7 +171,7 @@ def __init__(self, **kwargs):
 
         # Within the device="meta" context, tensors that are created do not carry data.
         # They possess all other metadata a tensor carries such as size, stride, requires_grad.
-        with torch.device("cpu"):
+        with torch.device("meta"):
             self.model_ = Transformer(model_args)
 
         if "int8" in str(checkpoint_path):