You need to set graph.owning_module on the graph module. This is because for call_module nodes, we need to fetch the original module object instance to check if it is set as pure (i.e. if it has no side effects) before removing it.
BTW, this is done usually by default when calling torch.fx.symbolic_trace(). Is there a reason you don’t want to call that helper instead?