amp opt_level=O0/O2/O3 results different after calling O1 in the same program

Hello,
I'm testing the speedup with amp with matmul operation on tesla v100(gcp). This is my code

class Matmul(torch.nn.Module):
    def __init__(self, n):
        super(Matmul, self).__init__()
        self.a = torch.randn(n, n).cuda()
        self.b = torch.randn(n, n).cuda()
    
    def forward(self):
        return torch.matmul(self.a, self.b)
    
def get_speed(n, num, use_amp=False, amp_opt_level='O1'):
    matmul = Matmul(n).cuda()
    if use_amp:
        from apex import amp
        matmul = amp.initialize(matmul, opt_level=amp_opt_level)
    else:
        amp = None
    start_time = time.time()
    for i in range(num):
        matmul()
    print('Finished in {:.3f} s.'.format(time.time() - start_time))
    matmul = None

I ran all these in one jupyter notebook: when I use opt_level=O0/O2/O3 before O1 I got

(no speedup)

I got 4x speedup with O1:

I got the same results when running them separately in the terminal.

However, when I use opt_level=O0/O2/O3 after O1 their speeds changed:

Is there any reason for that? Which result is correct?

Thank you.

Start with the provided get_speed benchmark and the amp.initialize entry point, running the opt_level cases in the reported orders in both the Jupyter notebook and terminal. Compare the timing results and determine which behavior is expected, documenting the cause and the correct result.