diff --git a/denoising_diffusion_pytorch/classifier_free_guidance.py b/denoising_diffusion_pytorch/classifier_free_guidance.py
index bf88ebfe4..6970ccec2 100644
--- a/denoising_diffusion_pytorch/classifier_free_guidance.py
+++ b/denoising_diffusion_pytorch/classifier_free_guidance.py
@@ -9,7 +9,7 @@
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from einops import rearrange, reduce, repeat
 from einops.layers.torch import Rearrange
@@ -731,7 +731,7 @@ def interpolate(self, x1, x2, classes, t = None, lam = 0.5):
 
         return img
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py
index 9ceccdbd8..df1940bf9 100644
--- a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py
+++ b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py
@@ -3,7 +3,7 @@
 from torch import sqrt
 from torch import nn, einsum
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.special import expm1
 
 from tqdm import tqdm
@@ -234,7 +234,7 @@ def sample(self, batch_size = 16):
 
     # training related functions - noise prediction
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index 3039e79f4..9b9d889cc 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -10,7 +10,7 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.utils.data import Dataset, DataLoader
 
 from torch.optim import Adam
@@ -772,7 +772,7 @@ def noise_assignment(self, x_start, noise):
         _, assign = linear_sum_assignment(dist.cpu())
         return torch.from_numpy(assign).to(dist.device)
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, t, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
index e94dde819..8f5c55ee2 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
@@ -9,7 +9,7 @@
 from torch import nn, einsum, Tensor
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.optim import Adam
 from torch.utils.data import Dataset, DataLoader
 
@@ -660,7 +660,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/guided_diffusion.py b/denoising_diffusion_pytorch/guided_diffusion.py
index e8856fed1..6e598885a 100644
--- a/denoising_diffusion_pytorch/guided_diffusion.py
+++ b/denoising_diffusion_pytorch/guided_diffusion.py
@@ -9,7 +9,7 @@
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.utils.data import Dataset, DataLoader
 
 from torch.optim import Adam
@@ -709,7 +709,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/repaint.py b/denoising_diffusion_pytorch/repaint.py
index 475ac0930..147564842 100644
--- a/denoising_diffusion_pytorch/repaint.py
+++ b/denoising_diffusion_pytorch/repaint.py
@@ -10,7 +10,7 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.utils.data import Dataset, DataLoader
 
 from torch.optim import Adam
@@ -815,7 +815,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, t, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/simple_diffusion.py b/denoising_diffusion_pytorch/simple_diffusion.py
index 6be766ca7..2b5b1b378 100644
--- a/denoising_diffusion_pytorch/simple_diffusion.py
+++ b/denoising_diffusion_pytorch/simple_diffusion.py
@@ -6,7 +6,7 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.special import expm1
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from tqdm import tqdm
 from einops import rearrange, repeat, reduce, pack, unpack
@@ -651,7 +651,7 @@ def sample(self, batch_size = 16):
 
     # training related functions - noise prediction
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py
index d841a5aa9..eedc77380 100644
--- a/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py
+++ b/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py
@@ -4,7 +4,7 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.special import expm1
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from tqdm import tqdm
 from einops import rearrange, repeat, reduce
@@ -150,7 +150,7 @@ def sample(self, batch_size = 16):
 
     # training related functions - noise prediction
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def q_sample(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py
index 37f038a14..ed6ed8969 100644
--- a/denoising_diffusion_pytorch/version.py
+++ b/denoising_diffusion_pytorch/version.py
@@ -1 +1 @@
-__version__ = '2.0.16'
+__version__ = '2.0.17'