Merge pull request #113 from laekov/compat-fix

Fix smart schedule with older PyTorch

fmoe/fastermoe/schedule.py

@@ -33,8 +33,13 @@ class MoEForward(Function):
             x = x.data
             with torch.enable_grad():
                 x.requires_grad = True
-                # To skip torch autograd's version check.
-                with torch.autograd.graph.saved_tensors_hooks(nothing, nothing):
+                try:
+                    # To skip torch autograd's version check.
+                    with torch.autograd.graph.saved_tensors_hooks(nothing, nothing):
+                        y0 = expert_fn(x, torch.tensor([x.shape[0]], dtype=torch.int64))
+                except Exception as e:
+                    # Ignore the error and fall back for compatibility to older
+                    # versions of PyTorch
                     y0 = expert_fn(x, torch.tensor([x.shape[0]], dtype=torch.int64))
             ctx.gibs[idx] = x
             ctx.gobs[idx] = y0