title: How to get the device type of a pytorch module conveniently?
body: <p>I have to stack some my own layers on different kinds of pytorch models with different devices.</p>
  
  <p>E.g. <strong>A</strong> is a <strong>cuda</strong> model and <strong>B</strong> is a <strong>cpu</strong> model (but I don't know it before I get the device type). Then the new models are <strong>C</strong> and <strong>D</strong> respectively, where</p>
  
  <pre class="lang-py prettyprint-override"><code>class NewModule(torch.nn.Module):
      def __init__(self, base):
          super(NewModule, self).__init__()
          self.base = base
          self.extra = my_layer() # e.g. torch.nn.Linear()
  
      def forward(self,x):
          y = self.base(x)
          z = self.extra(y)
          return z
  
  ...
  
  C = NewModule(A) # cuda
  D = NewModule(B) # cpu
  </code></pre>
  
  <p>However I must move <code>base</code> and <code>extra</code> to the <strong>same</strong> device, i.e. <code>base</code> and <code>extra</code> of C are <strong>cuda</strong> models and D's are <strong>cpu</strong> models. So I tried this <code>__inin__</code>:</p>
  
  <pre class="lang-py prettyprint-override"><code>def __init__(self, base):
      super(NewModule, self).__init__()
      self.base = base
      self.extra = my_layer().to(base.device)
  </code></pre>
  
  <p>Unfortunately, there's no attribute <code>device</code> in <code>torch.nn.Module</code>(raise <code>AttributeError</code>).</p>
  
  <p>What should I do to get the device type of <code>base</code>? Or any other method to make <code>base</code> and <code>extra</code> to be on the same device automaticly even the structure of <code>base</code> is unspecific?</p>