How can I move the dataloader from CPU to CUDA before starting training loop.

[DeepNeuralDog]

Hi @mrdbourke
I want to move the dataloader from CPU to CUDA before starting training loop**.(For dataloader: num_worker=2, pin_memory=True)
If I move the the input every time (cpu to gpu) for each epoch, my training time gets too much high(it is really really slow)

So I coded like this:

for train_features, train_labels in train_dataloader:
    train_features=train_features.to(device)
    train_labels=train_labels.to(device)
print('done')
for test_features, test_labels in test_dataloader:
    test_features=test_features.to(device)
    test_labels=test_labels.to(device)
print('done')


#BUT THIS METHOD DOES NOT WORK


in_channel=3
hidden= 32
output= len(class_names)

# MODEL INSTANCE
model_2= CNN(
    in_channel = in_channel,
    hidden = hidden,
    out_shape = output
).to(device)

# SET LOSS FN & OPTIMIZER

loss_fn= nn.CrossEntropyLoss()
optimizer= torch.optim.SGD(params=model_2.parameters(),
                            lr= 0.001
                            )



def accuracy_fn(y_true, y_pred):
  correct = torch.eq(y_true, y_pred).sum().item()
  acc = (correct/len(y_pred)) * 100
  return acc




epoch= 20
for i in tqdm(range(epoch)):
    # TRAINING
    train_loss=0
    train_acc=0
    for batch, (X, y) in enumerate(train_dataloader):
        # X,y= X.to(device),y.to(device) ## this line of code is taking most of the time.
        model_2.train()  

        y_pred= model_2(X) 

        loss = loss_fn(y_pred,y.to(device)) 
        train_loss+=loss 
        train_acc += accuracy_fn(y, y_pred.argmax(dim=1))

        optimizer.zero_grad() 

        loss.backward() 

        optimizer.step() 


    if i==0:
        print('Length of train dataloader',len(train_dataloader))
    train_loss /= len(train_dataloader)
    train_acc /= len(train_dataloader)

    #TESTING
    test_loss,test_acc = 0,0

    model_2.eval()
    with torch.inference_mode():


        for X_test, y_test in test_dataloader:
            # X_test,y_test= X_test.to(device),y_test.to(device)

            test_pred = model_2(X_test) 

            test_loss += loss_fn(test_pred,y_test) 

            test_acc += accuracy_fn(y_test, test_pred.argmax(dim=1))  

        test_loss /= len(test_dataloader)

        test_acc /= len(test_dataloader)

        print(f"Epoch: {i}| TrainLoss= {train_loss:4f}|TestLoss= {test_loss:4f}|TrainAccuracy= {train_acc:4f}|TestAccuracy = {test_acc:4f}")

Is there any easy fix?