ACE-Step Custom commited on
Commit
a11c404
·
1 Parent(s): 23dc982

Fix meta tensor: add device_map=None and use torch_dtype parameter

Browse files
Files changed (1) hide show
  1. acestep/handler.py +2 -1
acestep/handler.py CHANGED
@@ -492,7 +492,8 @@ class AceStepHandler:
492
  acestep_v15_checkpoint_path,
493
  trust_remote_code=True,
494
  attn_implementation=candidate,
495
- dtype=self.dtype,
 
496
  low_cpu_mem_usage=False, # Disable meta device for custom model with assertions
497
  )
498
  attn_implementation = candidate
 
492
  acestep_v15_checkpoint_path,
493
  trust_remote_code=True,
494
  attn_implementation=candidate,
495
+ torch_dtype=self.dtype,
496
+ device_map=None, # Disable automatic device mapping
497
  low_cpu_mem_usage=False, # Disable meta device for custom model with assertions
498
  )
499
  attn_implementation = candidate