Update on "[FSDP] Allow to use TorchDispatch with FSDP"

fegin · fegin · commit b5947416ab89 · 2022-11-03T19:30:05.000Z
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.

[ghstack-poisoned]
diff --git a/torch/distributed/fsdp/_utils.py b/torch/distributed/fsdp/_utils.py
@@ -9,6 +9,7 @@
     _is_namedtuple,
 )
 from torch.nn.utils.rnn import PackedSequence
+from torch.utils._mode_utils import no_dispatch
 
 
 def _contains_batchnorm(module):

Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@`
`9`	`9`	`_is_namedtuple,`
`10`	`10`	`)`
`11`	`11`	`from torch.nn.utils.rnn import PackedSequence`
	`12`	`+from torch.utils._mode_utils import no_dispatch`
`12`	`13`
`13`	`14`
`14`	`15`	`def _contains_batchnorm(module):`