Hi, I have a similar issue. I am trying to train a simple CNN network using my gpu (gtx 1070). Here is the code:
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pandas as pd 
from torchvision import datasets, transforms
from torchvision.utils import make_grid
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import time
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
## Create transform method 
transform = transforms.ToTensor()

#Download test and train data 
train_data = datasets.MNIST(root='C:\\Users\\aybars\\Google Drive\\pytorch',train=True,download=True,transform=transform)
test_data =datasets.MNIST(root='C:\\Users\\aybars\\Google Drive\\pytorch',train=False,download=True,transform=transform)
 
#create loader object for batch loading

train_loader = DataLoader(train_data,batch_size=10,shuffle=True,pin_memory=True)
test_loader =DataLoader(test_data,batch_size=10,shuffle=True,pin_memory=True)

#create conv layers 
conv1 = nn.Conv2d(1,6,3,1)
conv2 = nn.Conv2d(6,16,3,1)

class ConvolutionalNeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1,6,3,1)
        self.conv2 = nn.Conv2d(6,16,3,1)
        self.fc1 = nn.Linear(5*5*16,120)
        self.fc2 = nn.Linear(120,84)
        self.fc3 = nn.Linear(84,10)
        
    def forward(self,X):
        X = F.relu(self.conv1(X))
        X = F.max_pool2d(X,2,2)
        X = F.relu(self.conv2(X))
        X = F.max_pool2d(X,2,2)
        X = X.view(-1,5*5*16)
        X = F.relu(self.fc1(X))
        X = F.relu(self.fc2(X))
        X = F.log_softmax(self.fc3(X),dim=1)
        return X

torch.manual_seed(42)
model = ConvolutionalNeuralNetwork()
model = model.to(device)
print (model)


criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.001)


epoch = 2
train_losses = []
test_losses= []
train_correct=[]
test_correct =[]
t = time.time()
for i in range(epoch):
    trn_correct = 0
    tst_correct = 0
    for b,(X_train,y_train) in enumerate(train_loader):
        X_train=X_train.to(device)
        y_train=y_train.to(device)
        y_pred = model.forward(X_train)
        loss = criterion(y_pred,y_train)
        trn_correct+=(torch.max(y_pred,1)[1]==y_train).sum()
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if b%600==0:
            print(f"Epoch : {i} BATCH: {b} LOSS: {loss.item()}")

    train_losses.append(loss)
    train_correct.append(trn_correct)
    print(f'train_accuracy : {trn_correct.item()/len(train_data)}')
elapsed_time = time.time()-t
print(f'Duration : {elapsed_time}')

it takes roughly 150 seconds. However when I change
device = torch.device("cuda:0" if use_cuda else "cpu") 

to
device = torch.device("cpu" if use_cuda else "cpu") 

and setting pin_memory parameters to “False”
without changing anything else it takes 86 seconds.