Device pinned to tracing device in torch.jit.trace for torchscript

Misgana: <p>I am trying to load a torchscript on a different device from where it is traced, e.g cuda:3. The problem is the Device is a constant, and it is pinned to the tracing device (cuda:0). The map_location of <code>torch.jit.load</code> doesn’t help either.  Is there a way either a) to modify the traced graph and unpin the device constant or b) parameterize the device while tracing, so it is not pinned to a specific device.</p>
<pre><code class="lang-auto">graph(%self.1 : __torch__.mmedit.models.inpaintors.two_stage.TwoStageInpaintor,
      %masked_img.1 : Tensor,
      %mask.1 : Tensor):
  %11 : NoneType = prim::Constant()
  %10 : bool = prim::Constant[value=0]() # /mmediting/mmedit/models/inpaintors/two_stage.py:76:0
  %43 : Device = prim::Constant[value="cuda:0"]()
  %6 : int = prim::Constant[value=6]() # /mmediting/mmedit/models/inpaintors/two_stage.py:76:0
  %7 : int = prim::Constant[value=0]() # /mmediting/mmedit/models/inpaintors/two_stage.py:76:0
  %17 : int = prim::Constant[value=1]() # /mmediting/mmedit/models/inpaintors/two_stage.py:77:0
  %26 : float = prim::Constant[value=1.]() # /opt/conda/lib/python3.7/site-packages/torch/_tensor.py:544:0
  %4 : __torch__.mmedit.models.backbones.encoder_decoders.two_stage_encoder_decoder.DeepFillEncoderDecoder = prim::GetAttr[name="generator"](%self.1)
  %tmp_ones.1 : Tensor = aten::ones_like(%mask.1, %6, %7, %43, %10, %11) # /mmediting/mmedit/models/inpaintors/two_stage.py:76:0
  %16 : Tensor[] = prim::ListConstruct(%masked_img.1, %tmp_ones.1, %mask.1)
  %x.1 : Tensor = aten::cat(%16, %17) # /mmediting/mmedit/models/inpaintors/two_stage.py:77:0
  %21 : Tensor = prim::CallMethod[name="forward"](%4, %x.1) # :0:0
  %23 : Tensor = aten::mul(%21, %mask.1) # /mmediting/mmedit/models/inpaintors/two_stage.py:81:0
  %28 : Tensor = aten::rsub(%mask.1, %26, %17) # /opt/conda/lib/python3.7/site-packages/torch/_tensor.py:544:0
  %29 : Tensor = aten::mul(%masked_img.1, %28) # /mmediting/mmedit/models/inpaintors/two_stage.py:81:0
  %33 : Tensor = aten::add(%23, %29, %17) # /mmediting/mmedit/models/inpaintors/two_stage.py:81:0
  return (%33)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cuda:0! (when checking arugment for argument tensors in method wrapper__cat)


</code></pre>

azhanmohammed: <p>Maybe you can try the below given script and check:</p>
<pre><code class="lang-python">import torch
import io

# Load ScriptModule from io.BytesIO object
with open('scriptmodule.pt', 'rb') as f:
    buffer = io.BytesIO(f.read())

# Load all tensors to the original device
torch.jit.load(buffer)

# Load all tensors onto CPU, using a device
buffer.seek(0)
torch.jit.load(buffer, map_location=torch.device('cpu'))
</code></pre>
<p>Hopefully this should work.</p>

JuanFMontesinos: <p>I know very little about torch trace.<br>
My guess is you are using ops which trace the device. For example, slicing <code>tensor[1:7]</code> keeps track of the device. You should replace such operators by pytorch ones.</p>

Misgana: <p><a class="mention" href="/u/azhanmohammed">@azhanmohammed</a> Thanks for the reply!</p>
<p>The code snippet is doing the same as this <code>torch.load('torchscript.pt', map_location=torch.device('cpu'))</code><br>
and so unfortunately the above error still persists.</p>
<p>Maybe the best way is to modify the graph node i.e this constant <code>%43 : Device = prim::Constant[value="cuda:0"]()</code>  but I haven’t figured out a neat way to do it without it breaking</p>
<p><a class="mention" href="/u/juanfmontesinos">@JuanFMontesinos</a>  could you explain which tensor to slice?</p>

JuanFMontesinos: <p>The slicing is just an example of an of carried out in python that leads to a hardcoded tensor in jit.<br>
In your case it seems you are defining a constant.</p>
<pre><code class="lang-auto">  %10 : bool = prim::Constant[value=0]() # /mmediting/mmedit/models/inpaintors/two_stage.py:76:0
  %43 : Device = prim::Constant[value="cuda:0"]()
</code></pre>
<p>I don’t really know how it was created but maybe you can just make a pytorch buffer for it and it will work.</p>

purrfegt: <p>I’m still having the same issue here and can’t change the original code so I only have the traced models.<br>
I also don’t understand why map_location doesn’t work.<br>
Did anyone find a fix?</p>