Update on "Add autograd hook for python rpc call"

zhaojuanmao · zhaojuanmao · commit 5994e17eacd0 · 2019-10-17T16:21:35.000-07:00
1. currently if autograd context is valid, even tensors do not require grads and grads function are not attached. it still send rpc with autograd meta. This is not ideal. This diff makes some change to make sure rpc with autograd meta is sent only if autograd context is valid and tensors require grads 2. meanwhile create a utiliy to attach autograd info and functions as needed 3. add autograd send/recv functions for python rpc call 4. make changes to support nested python rpc calls 5. disallow nested dist autograd context (was landed in #27022) Differential Revision: [D17819153](https://our.internmc.facebook.com/intern/diff/D17819153/) [ghstack-poisoned]
diff --git a/test/dist_autograd_test.py b/test/dist_autograd_test.py
@@ -3,7 +3,6 @@
 import time
 import unittest
 
-from enum import IntEnum
 import torch
 import torch.distributed.autograd as dist_autograd
 import torch.distributed.rpc as rpc
