From 25063950fc72acaa9740ba936379b36df1c5e456 Mon Sep 17 00:00:00 2001
From: Janne Hellsten <jhellsten@nvidia.com>
Date: Fri, 19 Feb 2021 18:21:50 +0200
Subject: [PATCH] Print better error message for the dreaded upfirdn2d_plugin
 problem

Print full traceback when custom extension build fails.

Also allow pytorch 1.9 so that this runs against pytorch upstream
devel builds.

issues #2, #28, #35, #37, #39
---
 torch_utils/ops/bias_act.py            | 4 ++--
 torch_utils/ops/conv2d_gradfix.py      | 2 +-
 torch_utils/ops/grid_sample_gradfix.py | 2 +-
 torch_utils/ops/upfirdn2d.py           | 4 ++--
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/torch_utils/ops/bias_act.py b/torch_utils/ops/bias_act.py
index b092c7ffe..4bcb409a8 100755
--- a/torch_utils/ops/bias_act.py
+++ b/torch_utils/ops/bias_act.py
@@ -9,11 +9,11 @@
 """Custom PyTorch ops for efficient bias and activation."""
 
 import os
-import sys
 import warnings
 import numpy as np
 import torch
 import dnnlib
+import traceback
 
 from .. import custom_ops
 from .. import misc
@@ -47,7 +47,7 @@ def _init():
         try:
             _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
         except:
-            warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
+            warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
     return _plugin is not None
 
 #----------------------------------------------------------------------------
diff --git a/torch_utils/ops/conv2d_gradfix.py b/torch_utils/ops/conv2d_gradfix.py
index 4b6a9e958..e95e10d0b 100755
--- a/torch_utils/ops/conv2d_gradfix.py
+++ b/torch_utils/ops/conv2d_gradfix.py
@@ -50,7 +50,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.']):
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
         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/torch_utils/ops/grid_sample_gradfix.py b/torch_utils/ops/grid_sample_gradfix.py
index 1477be027..ca6b3413e 100755
--- a/torch_utils/ops/grid_sample_gradfix.py
+++ b/torch_utils/ops/grid_sample_gradfix.py
@@ -34,7 +34,7 @@ def grid_sample(input, grid):
 def _should_use_custom_op():
     if not enabled:
         return False
-    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.']):
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
         return True
     warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
     return False
diff --git a/torch_utils/ops/upfirdn2d.py b/torch_utils/ops/upfirdn2d.py
index f768b2ca3..ceeac2b98 100755
--- a/torch_utils/ops/upfirdn2d.py
+++ b/torch_utils/ops/upfirdn2d.py
@@ -9,10 +9,10 @@
 """Custom PyTorch ops for efficient resampling of 2D images."""
 
 import os
-import sys
 import warnings
 import numpy as np
 import torch
+import traceback
 
 from .. import custom_ops
 from .. import misc
@@ -31,7 +31,7 @@ def _init():
         try:
             _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
         except:
-            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
+            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
     return _plugin is not None
 
 def _parse_scaling(scaling):