This is the code for computing the main NCE-based loss, which seems most strongly affected by changes from 0.4.1 to 1.0.0. In 1.0.0, if I explicitly cast the input to F.log_softmax to torch.float64 and then cast back to torch.float32 before returining the loss, it changes performance compared to using torch.float32 throughout the computation. In 0.4.1 this casting seems to have no effect.
def compute_loss_glb2lcl_nce(rkhs_glb, rkhs_lcl, stabilize=True):
    # compute number of spatial locations for local features
    n_batch = int(rkhs_glb.size(0))
    n_rkhs = int(rkhs_glb.size(1))
    n_locs = int(rkhs_lcl.size(2) * rkhs_lcl.size(3))
    mask_mat = torch.eye(n_batch).unsqueeze(dim=2).cuda()              # (n_batch, n_batch, 1)
    # compute info cost for global features -> local features
    rkhs_pos = rkhs_lcl.reshape(n_batch, n_rkhs, -1).permute(0, 2, 1)  # (n_batch, n_locs, n_rkhs)
    rkhs_lcl_flat = rkhs_pos.reshape(-1, n_rkhs)                       # (n_batch * n_locs, n_rkhs)
    # compute scores on positive samples only
    pred_pos = torch.matmul(rkhs_pos, rkhs_glb.unsqueeze(dim=2))       # (n_batch, n_locs, 1)
    # compute scores on all samples
    all_scores = torch.mm(rkhs_glb, rkhs_lcl_flat.t())                 # (n_batch, n_batch*n_locs)
    all_scores = all_scores.reshape(n_batch, n_batch, n_locs)          # (n_batch, n_batch, n_locs)
    if stabilize:
        lgt_reg = 1e-3 * (all_scores**2.).mean()
    else:
        lgt_reg = 0. * (all_scores**2.).mean()
    # ...
    pred_neg = all_scores.reshape(n_batch, -1).unsqueeze(dim=1)        # (n_batch, 1, n_batch*n_locs)
    pred_neg = pred_neg.expand(-1, n_locs, -1)                         # (n_batch, n_locs, n_batch*n_locs)
    mask_mat = 1. - mask_mat.expand(-1, -1, n_locs)                    # (n_batch, n_batch, n_locs)
    mask_mat = mask_mat.reshape(n_batch, -1).unsqueeze(dim=1)          # (n_batch, 1, n_batch*n_locs)
    mask_mat = mask_mat.expand(-1, n_locs, -1)                         # (n_batch, n_locs, n_batch*n_locs)
    # do not include scores for negative sampler buffer (none was given)
    pred_msk = torch.cat([torch.ones_like(pred_pos), mask_mat], dim=2)  # (n_batch, n_locs, 1 + n_batch*n_locs)
    pred_lgt = torch.cat([pred_pos, pred_neg], dim=2)                   # (n_batch, n_locs, 1 + n_batch*n_locs)
    pred_lgt = (pred_msk * pred_lgt) + (10. * (pred_msk - 1.))          # (n_batch, n_locs, 1 + n_batch*n_locs)
    if stabilize:
        pred_lgt = tanh_clip(pred_lgt, 10.)
    # compute log softmax for positive samples
    pred_nll = -F.log_softmax(pred_lgt, dim=2)                         # (n_batch, n_locs, 1 + n_neg)
    loss_glb2lcl = pred_nll[:, :, 0].mean()                            # scalar :-)
    loss_glb2lcl = loss_glb2lcl + lgt_reg
    return loss_glb2lcl, None