Nit

prune.py

@@ -4,6 +4,8 @@ import re
 import torch
 from modeling_jamba import JambaForCausalLM
 
+output_dir = "/home/user/jamba-small"
+
 model = JambaForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map="cpu", torch_dtype=torch.bfloat16)
 
 def prune_and_copy_additional_layers(original_state_dict):
@@ -37,5 +39,5 @@ def prune_and_copy_additional_layers(original_state_dict):
 pruned_state_dict = prune_and_copy_additional_layers(model.state_dict())
 
 print("Saving weights...")
-torch.save(pruned_state_dict, '/scratch/nbrown9/jamba-small-v1.bin')
+torch.save(pruned_state_dict, output_dir)
 print("Done!")