Clean things up

This commit is contained in:
oobabooga 2023-01-16 16:35:45 -03:00
parent 3a99b2b030
commit 6456777b09
5 changed files with 16 additions and 23 deletions

View file

@ -17,7 +17,5 @@ model_name = path.name
print(f"Loading {model_name}...")
model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
print("Model loaded.")
print(f"Saving to torch-dumps/{model_name}.pt")
print(f"Model loaded.\nSaving to torch-dumps/{model_name}.pt")
torch.save(model, Path(f"torch-dumps/{model_name}.pt"))