Dataparallel pytorch share tensor

I’m having a model which I want to run in parallel on multiple gpus. The code works well, but I’m having a problem. I’m doing something that requires an input and an adjacency matrix (which is static). I’ve tried to save the adjacency matrix under the model but it only saves it on the first model only. I’ve tried to pass it as a parameter in the forward function but it splits it in half. Any ideas how I could do this?

Could you try to register this matrix as a buffer using self.register_buffer?

1 Like

Seems to be working. Thank you!

1 Like