Update: Adopt PR NVlabs/stylegan2-ada-pytorch#197 to fix `Update: Ado…

…pt PR NVlabs/stylegan2-ada-pytorch#197 to fix ` on PyTorch 1.10 and above.
bytedance · Jun 10, 2022 · 3eeeb81 · 3eeeb81
1 parent 388ed20
commit 3eeeb81
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/third_party/stylegan2_official_ops/conv2d_gradfix.py b/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

diff --git a/third_party/stylegan2_official_ops/grid_sample_gradfix.py b/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