Check for cuda is available for macos

This commit is contained in:
Bakhtiyor Sulaymonov 2025-02-26 22:18:18 +05:00
parent 9ab8f963c8
commit 9d3d4d784f

View File

@ -57,7 +57,7 @@ class WanT2V:
t5_cpu (`bool`, *optional*, defaults to False):
Whether to place T5 model on CPU. Only works without t5_fsdp.
"""
self.device = torch.device(f"cuda:{device_id}")
self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
self.config = config
self.rank = rank
self.t5_cpu = t5_cpu
@ -259,6 +259,7 @@ class WanT2V:
del sample_scheduler
if offload_model:
gc.collect()
if torch.cuda.is_available():
torch.cuda.synchronize()
if dist.is_initialized():
dist.barrier()