Cannot understand why window_pool is not on GPU. Entire module is on GPU, so wouldn’t window_pool be on GPU as well? Is gradient info copied when I do sliced tensor copying?
How do I make window_pool to be on GPU as well… Would it be correct just to do window_pool.to("cuda")?
def forward(self, input, lens, args):

        window_pool = torch.zeros([lens.shape[0], args.mem_dim*len(self.windows)])
    
        convolved = conv_model(x)[0].transpose(0, 1) 
        relu_convolved = F.relu(convolved)
        start = 0

        for i in range(lens.shape[0]):
                input_len = relu_convolved[start:start+lens[i]-window+1].shape[0]
                input = relu_convolved[start:start+lens[i]-window+1].permute(1,0).unsqueeze(0)

                start += lens[i].data.cpu().numpy()-window+1

                max_pool1d = F.max_pool1d(input, kernel_size=input_len)

                lo = (window-1)*args.mem_dim
                hi = lo + args.mem_dim

                print("max_pool1d.device ", max_pool1d.device) #GPU

                window_pool[i][lo:hi] = max_pool1d

        print(window_pool.device)  #CPU