I am trying to parallelize this network.

  
      github.com
  
  
    aimagelab/meshed-memory-transformer/blob/master/train.py#L182

    print("Building vocabulary")
    text_field.build_vocab(train_dataset, val_dataset, min_freq=5)
    pickle.dump(text_field.vocab, open('vocab_%s.pkl' % args.exp_name, 'wb'))
else:
    text_field.vocab = pickle.load(open('vocab_%s.pkl' % args.exp_name, 'rb'))


# Model and dataloaders
encoder = MemoryAugmentedEncoder(3, 0, attention_module=ScaledDotProductAttentionMemory,
                                 attention_module_kwargs={'m': args.m})
decoder = MeshedDecoder(len(text_field.vocab), 54, 3, text_field.vocab.stoi['<pad>'])
model = Transformer(text_field.vocab.stoi['<bos>'], encoder, decoder).to(device)


dict_dataset_train = train_dataset.image_dictionary({'image': image_field, 'text': RawField()})
ref_caps_train = list(train_dataset.text)
cider_train = Cider(PTBTokenizer.tokenize(ref_caps_train))
dict_dataset_val = val_dataset.image_dictionary({'image': image_field, 'text': RawField()})
dict_dataset_test = test_dataset.image_dictionary({'image': image_field, 'text': RawField()})




def lambda_lr(s):
    warm_up = args.warmup



  
  
    
    
  
  


After line 182 I just added
model = torch.nn.DataParallel(model)
I have tried Pytorch 1.1.0, 1.2.0, 1.4.0, and 1.5. (from source)