Update tools/llama/generate.py

tools/llama/generate.py

@@ -340,10 +340,7 @@ def encode_tokens(
 
 
 def load_model(checkpoint_path, device, precision, compile=False):
-    model: Union[NaiveTransformer, DualARTransformer] = BaseTransformer.from_pretrained(
-        checkpoint_path, load_weights=True, weights_only=True
-    )
-
+    
     weights = torch.load(checkpoint_path, weights_only=True)
     model = BaseTransformer.from_pretrained(weights)
     logger.info(f"Restored model from checkpoint")