address #328

Author: lucidrains

denoising_diffusion_pytorch/attend.py

@@ -60,7 +60,9 @@ def __init__(
 
         device_properties = torch.cuda.get_device_properties(torch.device('cuda'))
 
-        if device_properties.major == 8 and device_properties.minor == 0:
+        device_version = version.parse(f'{device_properties.major}.{device_properties.minor}')
+
+        if device_version > version.parse('8.0'):
             print_once('A100 GPU detected, using flash attention if input tensor is on cuda')
             self.cuda_config = AttentionConfig(True, False, False)
         else:

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '2.0.8'
+__version__ = '2.0.10'