What I am doing is in every batch I assign the tensor (named graph) to the module wrapped in DataParallel based on a criterion. This tensor is different for different batches and should not be parallelized, which is why I am trying to assign it directly to the model:
        for batch_idx, (x, y, graph, subject) in enumerate(self.train_loader):

            if model.module.subject != subject:

                model.module.subject = subject
                model.module.graph = graph      # where the assignment takes place

            output = model(x.to(self.device))
            target = torch.argmax(y, dim=1)

            optimizer.zero_grad()
            loss = F.nll_loss(output, target, weight=self.w)
            loss.backward()
            optimizer.step()

Given that I assign the tensor to the module wrapped by DataParallel, I would expect this tensor to be replicated in each GPU as well, which does not seem to be the case. This might have to do with the replication process:



 Andreas_Georgiou:

The model is replicated to each device after each iteration.


When exactly does the replication happen? And if this happens in every iteration, shouldn’t the model replicate this tensor as well?