So a bit of an update. I fixed the original issue without using retrain_graph=True. So no hacks involved.
I know that you previously said:

No, so the way it’s implemented is such that the mask keeps on getting replaced. We don’t keep on storing more and more masks with every iteration.

However, by the time I apply iterative pruning (approx. 100 times) an OOM still occurs. The message looks like this and seems more directly related to pruning:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/nfs_home/sohyongs/anaconda3/envs/fairseq/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
    fn(i, *args)
  File "/nfs_home/sohyongs/Workspace/fairseq/fairseq/distributed_utils.py", line 156, in distributed_main
    main(args, **kwargs)
  File "/nfs_home/sohyongs/Workspace/fairseq/fairseq_cli/train.py", line 121, in main
    valid_losses, should_stop = train(args, trainer, task, epoch_itr)
  File "/nfs_home/sohyongs/anaconda3/envs/fairseq/lib/python3.6/contextlib.py", line 52, in inner
    return func(*args, **kwds)
  File "/nfs_home/sohyongs/Workspace/fairseq/fairseq_cli/train.py", line 217, in train
    log_output = trainer.train_step(samples)
  File "/nfs_home/sohyongs/anaconda3/envs/fairseq/lib/python3.6/contextlib.py", line 52, in inner
    return func(*args, **kwds)
  File "/nfs_home/sohyongs/Workspace/fairseq/fairseq/trainer.py", line 475, in train_step
    amount=pruning_amount)
  File "/nfs_home/sohyongs/anaconda3/envs/fairseq/lib/python3.6/site-packages/torch/nn/utils/prune.py", line 1049, in global_unstructured
    final_mask = container.compute_mask(t, default_mask)
  File "/nfs_home/sohyongs/anaconda3/envs/fairseq/lib/python3.6/site-packages/torch/nn/utils/prune.py", line 391, in compute_mask
    mask = _combine_masks(method, t, default_mask)
  File "/nfs_home/sohyongs/anaconda3/envs/fairseq/lib/python3.6/site-packages/torch/nn/utils/prune.py", line 386, in _combine_masks
    new_mask[slc] = partial_mask.to(dtype=new_mask.dtype)
R

Maybe the _combine_masks part computation is requiring too much intermediate memory? Is it better practice to apply pruning on the CPU and then moving it to the device?
Due to it being more directly related to pruning, can you share any suspicions? @Michela Thanks!