Since the whole data has been transferred to GPU before creating dataset and train loader, i think it’s not having any loading time issues. pin_memory wouldn’t work for data already on GPU, if that’s what you’re suggesting. Also i have tried to minimize loss printing. Here are major parts of the code for clarity.
class Net(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        super(Net, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, H)
        self.linear3 = torch.nn.Linear(H, H)
        self.linear4 = torch.nn.Linear(H, H)
        self.linear5 = torch.nn.Linear(H, H)
        self.linear6 = torch.nn.Linear(H, H)
        self.linear7 = torch.nn.Linear(H, H)
        self.linear8 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        out = F.relu(self.linear1(x))
        out = F.relu(self.linear2(out))
        out = F.relu(self.linear3(out))
        out = F.relu(self.linear4(out))
        out = F.relu(self.linear5(out))
        out = F.relu(self.linear6(out))
        out = F.relu(self.linear7(out))
        y_pred = self.linear8(out)
        return y_pred

D_in, H, D_out = 3, 768, 1
model = Net(D_in, H, D_out)
criterion = nn.MSELoss(reduction='sum')
optimizer = Adam(model.parameters(), lr=5e-4)

device = torch.device('cuda')
model.to(device)

dataX = dataX.to(device)
dataY = dataY.to(device)
dataset = TensorDataset(dataX,dataY)
training_batches = 20
batch_size_train = int(len(dataX)/training_batches) +1
train_loader = DataLoader(dataset, batch_size=batch_size_train, shuffle=True)

start_time = time.time()
for epoch in range(201):
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        features, target = data
        optimizer.zero_grad()
        forward = model(features)
        loss = criterion(forward, target)
        if epoch % 100 == 0:
            running_loss += loss.item()
        loss.backward()
        optimizer.step()
    if epoch % 100 == 0:
        print('Epoch: {},    Training Loss: {:.2e}'.format(epoch, running_loss/training_batches))

elapsed = time.time() - start_time
print('GPU Time: {:.2f} min'.format(elapsed/60))