fffiloni commited on
Commit
f8f9211
·
verified ·
1 Parent(s): 3a6529b

Update skyreels_v2_infer/pipelines/image2video_pipeline.py

Browse files
skyreels_v2_infer/pipelines/image2video_pipeline.py CHANGED
@@ -117,7 +117,7 @@ class Image2VideoPipeline:
117
  )
118
 
119
  self.transformer.to(self.device)
120
- with torch.cuda.amp.autocast(dtype=self.transformer.dtype), torch.no_grad():
121
  self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift)
122
  timesteps = self.scheduler.timesteps
123
 
 
117
  )
118
 
119
  self.transformer.to(self.device)
120
+ with torch.amp.autocast("cuda", dtype=self.transformer.dtype), torch.no_grad():
121
  self.scheduler.set_timesteps(num_inference_steps, device=self.device, shift=shift)
122
  timesteps = self.scheduler.timesteps
123