# Buggy code
def binarize_with_mask(self, txt, prefix, suffix, leading_space, trailing_space):
    toks = self.binarize(
        prefix + leading_space + txt + trailing_space + suffix,
        append_eos=True,
    )
    mask = torch.zeros_like(toks, dtype=torch.uint8)
    mask_start = len(self.binarize(prefix))
    mask_size = len(self.binarize(leading_space + txt))
    mask[mask_start : mask_start + mask_size] = 1
    return toks, mask


# Traceback (NotImplementedError)
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | distributed init (rank 0): tcp://localhost:11739
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | distributed init (rank 1): tcp://localhost:11739
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | distributed init (rank 2): tcp://localhost:11739
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | initialized host ubuntu as rank 1
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | initialized host ubuntu as rank 2
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | distributed init (rank 3): tcp://localhost:11739
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | initialized host ubuntu as rank 3
2020-03-19 04:44:55 | INFO | fairseq.distributed_utils | initialized host ubuntu as rank 0
| dictionary: 50265 types
| dictionary: 50265 types
2020-03-19 04:45:03 | INFO | fairseq_cli.train | Namespace(activation_dropout=0.0, activation_fn='gelu', adam_betas='(0.9, 0.98)', adam_eps=1e-06, all_gather_list_size=16384, arch='roberta_base', attention_dropout=0.1, best_checkpoint_metric='accuracy', bpe='gpt2', broadcast_buffers=False, bucket_cap_mb=25, clip_norm=25, cpu=False, criterion='wsc', curriculum=0, data='WSC/', dataset_impl=None, ddp_backend='no_c10d', device_id=0, disable_validation=False, distributed_backend='nccl', distributed_init_method='tcp://localhost:11739', distributed_no_spawn=False, distributed_port=-1, distributed_rank=0, distributed_world_size=4, dropout=0.1, empty_cache_freq=0, encoder_attention_heads=12, encoder_embed_dim=768, encoder_ffn_embed_dim=3072, encoder_layerdrop=0, encoder_layers=12, encoder_layers_to_keep=None, end_learning_rate=0.0, fast_stat_sync=False, find_unused_parameters=False, fix_batches_to_gpus=False, fixed_validation_seed=None, force_anneal=None, fp16=True, fp16_init_scale=128, fp16_no_flatten_grads=False, fp16_scale_tolerance=0.0, fp16_scale_window=None, gpt2_encoder_json='https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json', gpt2_vocab_bpe='https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe', init_token=None, keep_best_checkpoints=-1, keep_interval_updates=-1, keep_last_epochs=-1, log_format='simple', log_interval=100, lr=[2e-05], lr_scheduler='polynomial_decay', max_epoch=0, max_positions=512, max_sentences=16, max_sentences_valid=16, max_tokens=None, max_tokens_valid=None, max_update=2000, maximize_best_checkpoint_metric=True, memory_efficient_fp16=False, min_loss_scale=0.0001, min_lr=-1, no_epoch_checkpoints=True, no_last_checkpoints=True, no_progress_bar=False, no_save=False, no_save_optimizer_state=True, num_workers=1, optimizer='adam', optimizer_overrides='{}', patience=-1, pooler_activation_fn='tanh', pooler_dropout=0.0, power=1.0, required_batch_size_multiple=8, reset_dataloader=True, reset_lr_scheduler=False, reset_meters=True, reset_optimizer=True, restore_file='/home/xgx/Commonsense/test/model/roberta.base/model.pt', save_dir='checkpoints', save_interval=1, save_interval_updates=0, save_predictions=None, seed=1, sentence_avg=False, skip_invalid_size_inputs_valid_test=False, task='wsc', tensorboard_logdir='', threshold_loss_scale=None, tokenizer=None, total_num_update=2000, train_subset='train', update_freq=[1], use_bmuf=False, use_old_adam=False, user_dir='/home/xgx/fairseq/examples/roberta/wsc', valid_subset='val', validate_interval=1, warmup_updates=250, weight_decay=0.01, wsc_cross_entropy=True, wsc_margin_alpha=1.0, wsc_margin_beta=0.0)
| dictionary: 50265 types
| dictionary: 50265 types
Traceback (most recent call last):
File "/home/xgx/miniconda3/bin/fairseq-train", line 11, in <module>
load_entry_point('fairseq', 'console_scripts', 'fairseq-train')()
File "/home/xgx/fairseq/fairseq_cli/train.py", line 317, in cli_main
nprocs=args.distributed_world_size,
File "/home/xgx/miniconda3/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 171, in spawn
while not spawn_context.join():
File "/home/xgx/miniconda3/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 118, in join
raise Exception(msg)
Exception:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
File "/home/xgx/miniconda3/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
fn(i, *args)
File "/home/xgx/fairseq/fairseq_cli/train.py", line 286, in distributed_main
main(args, init_distributed=True)
File "/home/xgx/fairseq/fairseq_cli/train.py", line 63, in main
criterion = task.build_criterion(args)
File "/home/xgx/fairseq/fairseq/tasks/fairseq_task.py", line 226, in build_criterion
return criterions.build_criterion(args, self)
File "/home/xgx/fairseq/fairseq/registry.py", line 41, in build_x
return builder(args, *extra_args, **extra_kwargs)
File "/home/xgx/fairseq/fairseq/criterions/fairseq_criterion.py", line 56, in build_criterion
'{}.build_criterion'.format(cls.__name__)
NotImplementedError: Unable to infer Criterion arguments, please implement WSCCriterion.build_criterion

# Fixed code
def binarize_with_mask(self, txt, prefix, suffix, leading_space, trailing_space):
    toks = self.binarize(
        prefix + leading_space + txt + trailing_space + suffix,
        append_eos=True,
    )
    mask = torch.zeros_like(toks, dtype=torch.bool)
    mask_start = len(self.binarize(prefix))
    mask_size = len(self.binarize(leading_space + txt))
    mask[mask_start : mask_start + mask_size] = 1
    return toks, mask