Merge pull request #4 from NVIDIA/amp_compat_fix

Fix compatibility checks for 18.04 container

apex/amp/compat.py

@@ -5,6 +5,10 @@ def variable_is_tensor():
     v = torch.autograd.Variable()
     return isinstance(v, torch.Tensor)
 
+def tensor_is_variable():
+    x = torch.Tensor()
+    return type(x) == torch.autograd.Variable
+
 # False for post-0.4
 def tensor_is_float_tensor():
     x = torch.Tensor()

apex/amp/lists/tensor_overrides.py

@@ -5,7 +5,7 @@ import importlib
 
 import torch
 
-if compat.variable_is_tensor():
+if compat.variable_is_tensor() and not compat.tensor_is_variable():
     MODULE = torch.Tensor
 else:
     MODULE = torch.autograd.Variable

apex/amp/wrap.py

@@ -8,8 +8,6 @@ import torch
 def cached_cast(mod, fn, cast_fn, handle,
                 try_caching=False, verbose=False):
     if not utils.has_func(mod, fn):
-        # Should happen only pre-0.4
-        assert not compat.variable_is_tensor()
         return
 
     orig_fn = utils.get_func(mod, fn)