I’m not sure if this is the reason but PyTorch recommends putting backward computation out of the autocast context.
Here’s a quote from the pytorch documentation 3.
autocast should wrap only the forward pass(es) of your network, including the loss computation(s). Backward passes under autocast are not recommended. Backward ops run in the same type that autocast used for corresponding forward ops.
In your case, it would be
    with torch.cuda.amp.autocast(): # CHANGE: Added amp autocast
        match_net = matching.MatchingNetwork(
            keep_prob=constants.keep_prob,
            batch_size=args.batch_size,
            num_channels=constants.num_channels,
            learning_rate=constants.lr,
            fce=constants.fce,
            image_size=args.image_dim,
            device=device,
        )

    scaler.scale(c_loss).backward() # on loss backprop
    scaler.step(optimiser)
    scaler.update()