remove cuda

cccclai · cccclai · commit a187cf8f664e · 2024-09-05T20:59:49.000-07:00
diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py
@@ -225,7 +225,7 @@ def calibrate_template(
         )
 
         eval_wrapper = EagerEvalWrapper(
-            model=prepared_module.to(device="cuda"),
+            model=prepared_module,
             tokenizer=tokenizer,
             max_seq_length=calibration_seq_length,
             use_kv_cache=self.use_kv_cache,

Original file line number	Diff line number	Diff line change
`@@ -225,7 +225,7 @@ def calibrate_template(`
`225`	`225`	`)`
`226`	`226`
`227`	`227`	`eval_wrapper = EagerEvalWrapper(`
`228`		`- model=prepared_module.to(device="cuda"),`
	`228`	`+ model=prepared_module,`
`229`	`229`	`tokenizer=tokenizer,`
`230`	`230`	`max_seq_length=calibration_seq_length,`
`231`	`231`	`use_kv_cache=self.use_kv_cache,`