python3 run.py with data_root=../../data/COCO num_gpus=2 num_nodes=1 task_mlm_itm whole_word_masking=True step200k per_gpu_batchsize=32 about vilt HOT 2 CLOSED

def set_task(pl_module):
    if not pl_module.training:
        pl_module.current_tasks = [
            k for k, v in pl_module.hparams.config["loss_names"].items() if v >= 1
        ]
        return

    if not check_non_acc_grad(pl_module):
        return

    sampling_pools = list()
    for k, v in pl_module.hparams.config["loss_names"].items():
        sampling_pools.extend([k] * int(v))

    g = torch.Generator()
    g.manual_seed(pl_module.global_step)
    idx = torch.randperm(len(sampling_pools), generator=g)[0]
    pl_module.current_tasks = sampling_pools[idx]