with self.assertRaisesRegex(RuntimeError, 'cannot contain a Node'):
traced_graph = MyTracer().trace(CallsModWithDict())
+ def test_module_deepcopy_edit_nodes(self):
+ class Foo(torch.nn.Module):
+ def forward(self, x):
+ return torch.relu(x)
+
+ traced1 = symbolic_trace(Foo())
+ copied = copy.deepcopy(traced1)
+
+ for node in copied.graph.nodes:
+ if node.target == torch.relu:
+ node.target = torch.neg
+
+ copied.recompile()
+ traced1.recompile()
+
+ x = torch.randn(15, 15)
+ torch.testing.assert_allclose(traced1(x), torch.relu(x))
+ torch.testing.assert_allclose(copied(x), torch.neg(x))
+
def test_direct_param_use(self):
class TransposeTest(torch.nn.Module):
def __init__(self):
def __deepcopy__(self, memo):
fake_mod = torch.nn.Module()
fake_mod.__dict__ = copy.deepcopy(self.__dict__)
- return GraphModule(fake_mod, self.graph)
+ return GraphModule(fake_mod, fake_mod.__dict__['_graph'])
def __copy__(self):
return GraphModule(self, self.graph)