//  [python]
def evaluate_VIC(train_loader, model, args, log):
    losses = AverageMeter()
    rank_calculate = RankCalculator()

    local_rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()

    # switch to train mode
    model.train()

    iters_per_epoch = len(train_loader)
    moco_m = args.moco_m

    VIC_rec = [{"Var": AverageMeter(), "InVar": AverageMeter(), "CoVar": AverageMeter(), "Norm": AverageMeter()} for _ in range(len(model.module.base_encoder.blocks))]

    with torch.no_grad():
        for i, (images, _) in enumerate(train_loader):
            if args.local_rank is not None:
                images[0] = images[0].cuda(args.local_rank, non_blocking=True)
                images[1] = images[1].cuda(args.local_rank, non_blocking=True)

            # compute output
            with torch.cuda.amp.autocast(True):
                model(images[0], images[1], moco_m, record_feature=True)
            features = model.module.features

            # aggregate features, compute VIC for each layer
            for key, item in features.items():
                for cnt in range(len(item)):
                    item[cnt] = gather_features(item[cnt], local_rank, world_size)

            if torch.distributed.get_rank() == 0:
                for layer_cnt in range(len(features["x1"])):
                    # cal VIC layer by layer
                    x = features["x1"][layer_cnt].mean(1)
                    y = features["x2"][layer_cnt].mean(1)

                    # compute norm
                    norm_x = torch.linalg.norm(x, ord=2, dim=-1)
                    norm_y = torch.linalg.norm(y, ord=2, dim=-1)
                    VIC_rec[layer_cnt]["Norm"].update(((norm_x + norm_y) / 2).mean().item())

                    # normalize
                    x_norm = x / norm_x.unsqueeze(-1)
                    y_norm = y / norm_y.unsqueeze(-1)

                    mse_loss = torch.nn.functional.mse_loss(x_norm, y_norm)
                    VIC_rec[layer_cnt]["InVar"].update(mse_loss.item())

                    x = x - x.mean(dim=0)
                    y = y - y.mean(dim=0)
                    norm_x = torch.linalg.norm(x, ord=2, dim=-1)
                    norm_y = torch.linalg.norm(y, ord=2, dim=-1)
                    # calculate the norm after remove mean
                    x = x / norm_x.unsqueeze(-1)
                    y = y / norm_y.unsqueeze(-1)

                    std_x = torch.sqrt(x.var(dim=0) + 0.0001).mean()
                    std_y = torch.sqrt(y.var(dim=0) + 0.0001).mean()
                    VIC_rec[layer_cnt]["Var"].update(((std_x + std_y) / 2).item())

                    batch_size = x.shape[0]
                    cov_x = (x.T @ x) / (batch_size - 1)
                    cov_y = (y.T @ y) / (batch_size - 1)
                    cov_loss = (off_diagonal(cov_x).pow_(2).sum() + off_diagonal(cov_y).pow_(2).sum()) / 2
                    VIC_rec[layer_cnt]["CoVar"].update(cov_loss.item())

            if i % 10 == 0:
                msg = "{}/{}, loss avg is {:.3f}".format(i, iters_per_epoch, losses.avg)
                log.info(msg)

    if torch.distributed.get_rank() == 0:
        VIC_result = [{key: rec[key].avg for key in rec} for rec in VIC_rec]
        np.save(os.path.join(log.path, "VIC_result"), VIC_result)