Question about pin_memory and Dataloaders

Anthony_Weaver: <p>I’m trying to use pin_memory and then later non_blocking to speed up transfers to the GPU and I am running into a problem where it seems like the data is not being transferred to the GPU.<br>
The error is:<br>
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor</p>
<p>Any help as to why/where the problem is would be helpful.  Several code snippets below provide details on my setup</p>
<p>My dataset is a custom dataset (custom PulsarData class) and I added a pin_memory method as follows:</p>
<pre><code>def pin_memory(self):
    self.ft_data = self.ft_data.pin_memory()
    self.dt_data = self.dt_data.pin_memory()
    self.labels = self.labels.pin_memory()

    return self
</code></pre>
<p>In addition this class has the following <strong>getitem</strong> method</p>
<pre><code>def __getitem__(self, index: int)-&gt; tuple:
    ft_data = np.empty((*self.ft_dim, self.n_channels))
    dt_data = np.empty((*self.dt_dim, self.n_channels))

    # Do some processing before passing observation to model 
    ft_data = s.detrend(np.nan_to_num(np.array(self.ft_data[index], dtype=np.float32).T))
    ft_data /= np.std(ft_data)
    ft_data -= np.median(ft_data)
    
    dt_data = np.nan_to_num(np.array(self.dt_data[index], dtype=np.float32))
    dt_data /= np.std(dt_data)
    dt_data -= np.median(dt_data)

    ft_data = np.reshape(ft_data, (self.n_channels, *self.ft_dim))
    dt_data = np.reshape(dt_data, (self.n_channels, *self.dt_dim))

    # Return data as PyTorch Tensor
    return torch.from_numpy(ft_data), torch.from_numpy(dt_data), torch.tensor(self.labels[index])
</code></pre>
<p>The dataloader in my training code is as follows:</p>
<pre><code>train_data = PulsarData(files=train_data_files)
train_data, validate_data = random_split(train_data, [0.85, 0.15])

tr_dataloader = DataLoader(train_data, batch_size=args.batch_size, pin_memory=True, shuffle=True)
v_dataloader = DataLoader(validate_data, batch_size=args.batch_size, pin_memory=True, shuffle=False)
</code></pre>
<p>I pass the dataloader to another function that performs a single training pass (NOTE: DEVICE is essentially a constant for cuda.  I use the same value to load the model to the GPU):</p>
<pre><code>def train_loop(dataloader: DataLoader, 
           model: nn.Module,
           data: str,
           loss_fn: _Loss, 
           optimizer: Optimizer,
           batch_size: int,
) -&gt; None:

size = len(dataloader.dataset)

# Set the model to training mode - important for batch normalization and dropout layers
model.train()

start_time = time.time()

for batch_idx, (freq_data, dm_data, labels) in enumerate(dataloader):
    batch_data = None

    # Load labels to device
    labels = labels.to(DEVICE, non_blocking=True)

    # Add some noise to freq data to help avoid overtraining
    if data == "freq":
        noise = torch.randn_like(freq_data) * .1
        batch_data = freq_data + noise
    elif data == "dm":
        batch_data = dm_data
    else:
        print(f"Invalid data type provided: {data}", flush=True)
        sys.exit(0)

    batch_data.to(DEVICE, non_blocking=True)
    predicted = model(batch_data)
</code></pre>

ptrblck: <aside class="quote no-group" data-username="Anthony_Weaver" data-post="1" data-topic="213280">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/anthony_weaver/48/72576_2.png" class="avatar"> Anthony_Weaver:</div>
<blockquote>
<pre><code class="lang-auto">    batch_data.to(DEVICE, non_blocking=True)
    predicted = model(batch_data)
</code></pre>
</blockquote>
</aside>
<p><code>batch_data</code> is not moved to <code>DEVICE</code> since you need to assign the output to a variable:</p>
<pre data-code-wrap="python"><code class="lang-python">batch_data = batch_data.to(...)
</code></pre>