Merge pull request #113 from laekov/compat-fix

Fix smart schedule with older PyTorch

fmoe/fastermoe/schedule.py

@@ -33,9 +33,14 @@ class MoEForward(Function):
             x = x.data
             with torch.enable_grad():
                 x.requires_grad = True
+                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
             y.copy_(y0)