address #328

lucidrains · Jun 21, 2024 · b83c401 · b83c401
1 parent 3bee377
commit b83c401
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/denoising_diffusion_pytorch/attend.py b/denoising_diffusion_pytorch/attend.py
@@ -60,7 +60,7 @@ def __init__(
 
         device_properties = torch.cuda.get_device_properties(torch.device('cuda'))
 
-        if device_properties.major == 8 and device_properties.minor == 0:
+        if device_properties.major >= 8 and device_properties.minor == 0:
             print_once('A100 GPU detected, using flash attention if input tensor is on cuda')
             self.cuda_config = AttentionConfig(True, False, False)
         else:

diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py
@@ -1 +1 @@
-__version__ = '2.0.8'
+__version__ = '2.0.9'