Repository: pytorch
Issue: optimize_for_mobile has an issue with constant operations at the end of a loop
Labels: oncall: mobile

### 🐛 Describe the bug

I've been working on a model that I want to optimize for mobile, but when optimizing the model for mobile, I ran into an issue with an operation at the end of a loop. I'll use the code example below to show the issue:

```
import torch
from torch import nn
from torch.utils.mobile_optimizer import optimize_for_mobile




class testModel(nn.Module):
    def __init__(self):
        super(testModel, self).__init__()
        
        # Input is 5, output is 5
        self.m = nn.ModuleList([nn.Linear(5, 5) for i in range(5)])
    
    
    
    # Forward takes noise as input and returns 5 outputs
    # from the network
    def forward(self, X):
        # the output is initially all ones
        Y = torch.ones((5))
        
        # array to hold the outputs
        out = []
        
        # Iterate 5 times to get 5 outputs
        for i in range(0, 5):
            # Send the inputs through the blocks
            for b in self.m:
                Y = b(Y)
            
            # Save the ith output
            out.append(Y[i])
            
            # Add a constant to the ith value of Y. This
            # is what breaks the optimized model
            Y[i] *= 2
        
        return torch.stack(out)



def main():
    # Create a new network
    model = testModel()
    
    # Trash input
    X = torch.zeros((5))
    
    # Get the network output
    Y = model(X)
    
    # Create torch script form of the model
    ts_model = torch.jit.trace(model, X)
    ts_model_mobile = optimize_for_mobile(ts_model)
    
    # Get the output from the torch script models
    Y_ts_model = ts_model(X)
    Y_ts_model_mobile = ts_model_mobile(X)
    
    print(ts_model_mobile.code)
    
    # What are the outputs from the model?
    print("Original: ", Y)
    print("Optimized: ", Y_ts_model)
    print("Mobile Optimized: ", Y_ts_model_mobile)
    
    # The output of the model should be the same as the
    # output from the optimized models
    assert torch.all(Y.eq(Y_ts_model)), "Torch script different from original model"
    assert torch.all(Y.eq(Y_ts_model_mobile)), "Mobile torch script different from original model"
    
    
    
if __name__=='__main__':
    main()
```

If you run the code, you would expect the output of the torch script model to be the same as the output of the original model, but the outputs of the models are slightly off. The original traced module works perfectly fine, but the mobile script runs into an issue where the output is slightly off for all values in the output vector (excluding the first). Taking a look at the mobile torch script code, I think I see what the issue is:

![image](https://user-images.githubusercontent.com/43501738/179577881-b57265ac-99d3-4428-aaf1-64c9a4e4be0b.png)

At the beginning of each loop, it applies the constant multiplication operation. The first operation is applied at _3 and _4. This operation should be applied after the loop instead of at the beginning.

Another issue (which I think is the main problem) is after the multiplication operation is performed on the Y tensor (at _3 and _4), the value is never used. As seen in _5, the _1 tensor is used which is the tensor before the constant operation is applied. Instead, the forward method at _5 should probably be using _4.

### Versions

torch: 1.12.0