Update wan/vace.py

wan/vace.py

@@ -15,7 +15,7 @@ from PIL import Image
 import torchvision.transforms.functional as TF
 import torch
 import torch.nn.functional as F
-import torch.cuda.amp as amp
+import torch.amp as amp
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from tqdm import tqdm
@@ -362,7 +362,7 @@ class WanVace(WanT2V):
         no_sync = getattr(self.model, 'no_sync', noop_no_sync)
 
         # evaluation mode
-        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+        with amp.autocast("cuda", dtype=self.param_dtype), torch.no_grad(), no_sync():
 
             if sample_solver == 'unipc':
                 sample_scheduler = FlowUniPCMultistepScheduler(
@@ -616,7 +616,7 @@ class WanVaceMP(WanVace):
                 no_sync = getattr(model, 'no_sync', noop_no_sync)
 
                 # evaluation mode
-                with amp.autocast(dtype=param_dtype), torch.no_grad(), no_sync():
+                with amp.autocast("cuda", dtype=param_dtype), torch.no_grad(), no_sync():
 
                     if sample_solver == 'unipc':
                         sample_scheduler = FlowUniPCMultistepScheduler(