Pytorch DataParallel with custom model
I want to train model with multiple gpu's. I'm using following code
model = load_model(path)
if torch.cuda.device_count() &gt; 1:
  print(&quot;Let's use&quot;, torch.cuda.device_count(), &quot;GPUs!&quot;)
  # dim = 0 [30, xxx] -&gt; [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)

model.to(device)

It works well except DataParallel doesn't contain functions from original model, is there a way around it? Thank you

The nn.Module passed to nn.DataParallel will end up being wrapped by the class to handle data parallelism. You can still access your model with the module attribute.
&gt;&gt;&gt; p_model = nn.DataParallel(model)
&gt;&gt;&gt; p_model.module # &lt;- model

For instance, to access your underlying model's quantize attribute, you would do:
&gt;&gt;&gt; p_model.module.quantize

