Tensor on different device

chagelo: <pre><code class="lang-python">    for i, data in enumerate(dataset):
        img = torch.squeeze(data['image'].to(device))
        
        save_path = data['save_path']
        
        h, w = img.shape
        
        num = list(range(0,int(w / 2 - 10))) + list(range(int(w / 2 + 10), w)) 
        # motion_events = np.random.randint(30, high=150)
        motion_events = 1
        k_space_lines = random.sample(num, motion_events)

        class AffineTransform(torch.nn.Module):
            def __init__(self, h, w, n) -&gt; None:
                super().__init__()
                self.n = n
                self.I = torch.nn.Parameter(torch.randn((1, 1, h, w), dtype=torch.float))
                self.X = torch.nn.Parameter(torch.randn(n, dtype=torch.float))
                self.Y = torch.nn.Parameter(torch.randn(n, dtype=torch.float))
                self.THETA = torch.nn.Parameter(torch.randn(n, dtype=torch.float))
                self.tm = torch.stack([torch.cos(self.THETA), torch.sin(-self.THETA), self.X, torch.sin(self.THETA), torch.cos(self.THETA), self.Y], dim=1).reshape(n, 1, 2, 3)
            def forward(self, i):
                self.tm = self.tm.to(device)
                grid = F.affine_grid(self.tm[i], torch.Size(self.I.size()), align_corners=False)
                return F.grid_sample(self.I, grid, align_corners=False).to(device)
                # return self.I
            def update(self):
                self.tm = torch.stack([torch.cos(self.THETA), torch.sin(-self.THETA), self.X, torch.sin(self.THETA), torch.cos(self.THETA), self.Y], dim=1).reshape(self.n, 1, 2, 3)

        stn = AffineTransform(h, w, motion_events).to(device)
        optimz = torch.optim.Adam(stn.parameters(), lr=0.001)

        print("processing: {}".format(save_path))
        for epoch in range(10000):
            
            corrupted_img = stn(0)

            lsf = torch.nn.L1Loss().to(device)
            optimz.zero_grad()

            diff = lsf(img, corrupted_img[0, 0])
            diff.backward()
            optimz.step()
            stn.update()
</code></pre>
<p>I debug find that all tensor in cuda, but it throw me a <code>RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:6 and cpu!</code> I don’t know why.</p>

srishti-git1110: <p>If you’ve found that all tensors are on the correct device, please also make sure that the model instance is on the correct device.</p>

chagelo: <p><code>stn = AffineTransform(h, w, motion_events).to(device)</code>, obviously.</p>

vdw: <p>I don’t think that <code>torch.nn.Parameter(...)</code> does suffice so that these tensor get moved to the device with <code>.to(device)</code> for the models.</p>
<p>You can wrap wrap all parameters in a <code>nn.ModuleList</code> or something, or change your code to something like</p>
<pre><code class="lang-auto">class AffineTransform(torch.nn.Module):
            def __init__(self, h, w, n, device='cpu') -&gt; None:
                super().__init__()
                self.n = n
                self.I = torch.nn.Parameter(torch.randn((1, 1, h, w), dtype=torch.float)).to(device)
                ...

stn = AffineTransform(h, w, motion_events, device=device).to(device)
</code></pre>
<p>Maybe there are more elegant solutions but it should work as I just had this issue myself.</p>

chagelo: <p>So, I try much more, and finally solve it. What I do is put <code>tm = torch.stack([torch.cos(self.THETA), torch.sin(-self.THETA), self.X, torch.sin(self.THETA), torch.cos(self.THETA), self.Y], dim=1).reshape(n, 1, 2, 3)</code> in <code>forward</code>, but i don’t know why <img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=12" title=":joy:" class="emoji" alt=":joy:" loading="lazy" width="20" height="20"></p>