Make backward compatible with older version of kornia

greentfrapp · May 1, 2021 · b8bbd42 · b8bbd42
1 parent 8a6f7bb
commit b8bbd42
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/lucent/optvis/transform.py b/lucent/optvis/transform.py
@@ -24,6 +24,7 @@
 
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+KORNIA_VERSION = kornia.__version__
 
 
 def jitter(d):
@@ -68,7 +69,10 @@ def inner(image_t):
         # kornia takes degrees
         alpha = _rads2angle(np.random.choice(angles), units)
         angle = torch.ones(b) * alpha
-        scale = torch.ones(b, 2)
+        if KORNIA_VERSION < '0.4.0':
+            scale = torch.ones(b)
+        else:
+            scale = torch.ones(b, 2)
         center = torch.ones(b, 2)
         center[..., 0] = (image_t.shape[3] - 1) / 2
         center[..., 1] = (image_t.shape[2] - 1) / 2