low cpu mem mode

model-tracing/main.py

@@ -91,9 +91,7 @@ results["commit"] = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode
 manual_seed(args.seed)
 
 dtype = torch.bfloat16
-low_cpu_mem_usage = (
-    "70b" in args.base_model_id.lower()
-)  # Enable low memory loading if "70b" is in model name
+low_cpu_mem_usage = True
 
 print(f"Low CPU Mem Usage Flag set to {low_cpu_mem_usage}")
 base_model = AutoModelForCausalLM.from_pretrained(