Update tools/llama/generate.py

tools/llama/generate.py

@@ -344,7 +344,8 @@ def load_model(checkpoint_path, device, precision, compile=False):
         checkpoint_path, load_weights=True, weights_only=True
     )
 
-    model = model.to(device=device, dtype=precision)
+    weights = torch.load(checkpoint_path, weights_only=True)
+    model = BaseTransformer.from_pretrained(weights)
     logger.info(f"Restored model from checkpoint")
 
     if isinstance(model, DualARTransformer):