Hi,
There is a sharing tensor (say A) in my model, which is computed at the begining of each training epoch.
I do not want A to be computed parallel, but the rest of my model runs in parallel.
In order to do this, I compute and register A to model instance every epoch, so my running process is like the code below:
def run_epoch(self):
    self.model.A  = self.compute_A()
    result = self.model(input) # model is instance of nn.DataParallel, and will use A in this calling

However I get the error below:
RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:253

Currently I think computing A at each GPU will solve this problem, but this method will consume more time, and I do not want to waste computing time. So how can I solve it elegantly?