AMP failing on specific line of code

[JohnGiorgi]

I am using amp with an opt_level="O1". It fails on a specific line of code, saying it expected argument #2 to be of type half but found type float

max_val = torch.max(pos_pairs, torch.max(neg_pairs, dim=1, keepdim=True)[0])

The line comes from here

Any idea why amp is failing to handle this line of code properly? I am using opt_level="01".