Update: Adopt PR NVlabs/stylegan2-ada-pytorch#197 to fix Update: Adopt PR NVlabs/stylegan2-ada-pytorch#197 to fix on PyTorch 1.10 and above.

Author: BrandoZhang

third_party/stylegan2_official_ops/conv2d_gradfix.py

@@ -29,6 +29,7 @@
 import warnings
 import contextlib
 import torch
+from distutils.version import LooseVersion
 
 enabled = True                     # Enable the custom op by setting this to true.
 weight_gradients_disabled = False  # Forcefully disable computation of gradients with respect to the weights.
@@ -61,7 +62,7 @@ def _should_use_custom_op(input):
         return False
     if input.device.type != 'cuda':
         return False
-    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
+    if LooseVersion(torch.__version__) >= LooseVersion('1.7.0'):
         return True
     warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
     return False

third_party/stylegan2_official_ops/grid_sample_gradfix.py

@@ -26,6 +26,7 @@
 
 import warnings
 import torch
+from distutils.version import LooseVersion
 
 #----------------------------------------------------------------------------
 
@@ -43,7 +44,7 @@ def grid_sample(input, grid, impl='cuda'):
 def _should_use_custom_op():
     if not enabled:
         return False
-    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
+    if LooseVersion(torch.__version__) >= LooseVersion('1.7.0'):
         return True
     warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
     return False