# \[Fatal Bug\] changed nn.Module.training does not trigger recompilation

**URL:** <https://dev-discuss.pytorch.org/t/fatal-bug-changed-nn-module-training-does-not-trigger-recompilation/1386>\
**Category:** compiler\
**Created:** [July 20, 2023, 8:53am UTC](https://dev-discuss.pytorch.org/t/fatal-bug-changed-nn-module-training-does-not-trigger-recompilation/1386 "2023-07-20T08:53:21Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![youkaichao](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/youkaichao/32/1278_2.png) [@youkaichao](https://dev-discuss.pytorch.org/u/youkaichao)\
**Post date:** [July 20, 2023, 8:53am UTC](https://dev-discuss.pytorch.org/t/fatal-bug-changed-nn-module-training-does-not-trigger-recompilation/1386/1 "2023-07-20T08:53:21Z")

</div>

### 🐛 Describe the bug

Here is the sample code:

```python
import torch
from torch import nn

from torch import _dynamo
from torch._functorch.aot_autograd import aot_module_simplified
import functorch
from functorch.compile import make_boxed_func

def graph_processing_pytorch(gm, example_inputs):
    # graph transform (graph optimization) for the captured graph in pytorch
    print("captured graph in pytorch")
    gm.print_readable()

def graph_processing_aot_forward(gm, example_inputs):
    # graph transform (graph optimization) for the captured graph in aot autograd forward graph
    print("captured graph in aot autograd forward")
    gm.print_readable()

def graph_processing_aot_backward(gm, example_inputs):
    # graph transform (graph optimization) for the captured graph in aot autograd backward graph
    print("captured graph in aot autograd backward")
    gm.print_readable()

def forward_compiler(gm, example_inputs):
    graph_processing_aot_forward(gm, example_inputs)
    return make_boxed_func(gm.forward)

def backward_compiler(gm, example_inputs):
    graph_processing_aot_backward(gm, example_inputs)
    return make_boxed_func(gm.forward)

def custom_backend(gm, example_inputs):
    graph_processing_pytorch(gm, example_inputs)
    return aot_module_simplified(
        gm, example_inputs,
        fw_compiler=forward_compiler,
        bw_compiler=backward_compiler
    )

class BackboneModel(nn.Module):

    def __init__ (self, *args, **kwargs) -> None:
        super(). __init__ (*args, **kwargs)
        self.drop1 = nn.Dropout(0.5)
        self.bn1 = nn.BatchNorm2d(6)

    def forward(self, x):
        x = self.bn1(self.drop1(x))
        return x

model = BackboneModel()

opt_model = torch.compile(model, backend=custom_backend)

input = torch.randn(64, 6, 32, 32)

output1 = opt_model(input)

# calling .eval in the whole model works
# opt_model.eval()

# calling .eval in sub model does not trigger re-compilation
opt_model.drop1.eval()
opt_model.bn1.eval()

output2 = opt_model(input)
output3 = opt_model(input)

print((output2 - output3).abs().max().item()) # huge number, nondeterministic!

```

It seems only the `.eval` call in the whole model works. The `training` flag in sub-models are not respected. As a result, after I call `eval` for all the sub-models, the output is still nondeterministic.

### Versions

pytorch 2.0.1

Note: I opened an issue at [[Fatal Bug] changed nn.Module.training does not trigger recompilation. · Issue #105653 · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/issues/105653) . Recently I opened several issues, but got no response. Therefore, I come here for dev support ☹