I am using a GeForce GTX 1060 6GB/PCIe/SSE2 and a AMD® Ryzen 5 1600 six-core processor × 12
I am trying to train a simple model on flower photos dataset
My code hangs my whole pc when training starts
Heres my code
import torch
from torchvision import datasets, transforms, models
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from train import train
from collections import OrderedDict
import matplotlib.pyplot as plt
import numpy as np
import helper
from fashion_mnist_conv import FeatureNet

model = models.resnet34(pretrained=True)
train_transform = transforms.Compose([transforms.RandomRotation(30),
                                transforms.RandomResizedCrop(224),
                                transforms.RandomHorizontalFlip(),
                                transforms.ToTensor(),
                                
                                ])


test_transform = transforms.Compose([transforms.Resize(124,124), transforms.ToTensor()])

testset = datasets.ImageFolder(root='flowers/test', transform=test_transform)
trainset = datasets.ImageFolder(root='flowers/train', transform=train_transform)


testloader = DataLoader(testset, batch_size=1, shuffle=True)
trainloader = DataLoader(trainset, batch_size=1, num_workers=1, shuffle=True)


image, label = next(iter(trainloader))
featnet = FeatureNet()
output = featnet(image)

#for param in model.parameters():
#    param.requires_grad = False

classifier = nn.Sequential(OrderedDict([
                ('fc1', nn.Linear(512, 300)),
                ('relu', nn.ReLU()),
                ('fc2', nn.Linear(300, 4)),
                ('output', nn.LogSoftmax(dim=1))
                ]))

model.fc = classifier
featnet = FeatureNet()

print("getting in cuda..")
criterion = nn.NLLLoss()
criterion2 = nn.CrossEntropyLoss()

train(10, featnet, criterion2, trainloader, testloader, cuda)


if you change the featnet with model(resnet34) same thing happens.
heres featnet
import torch
import torch.nn as nn


class FeatureNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3,8, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(8,20, kernel_size=5) 
        self.conv3 = nn.Conv2d(20, 30, kernel_size=5)
        self.dropout = nn.Dropout(p=0.5)

        self.mp = nn.MaxPool2d(1)
        

        self.fc = nn.Linear(1348320, 10)

    def forward(self, x):
        x = self.dropout(torch.relu(self.mp(self.conv1(x))))
        x = self.dropout(torch.relu(self.mp(self.conv2(x))))
        x = self.dropout(torch.relu(self.mp(self.conv3(x))))
        x = x.view(x.shape[0], -1)
        print(x.shape)

        x = self.fc(x) 

        return x


and heres my train.py
import torch
import torch.optim as optim

def test(model, testloader, criterion, cuda):
    equals = 0
    n = 0
    running_loss = 0
    with torch.no_grad():
        for x, y in testloader:
            if cuda:
               x,y = x.cuda(), y.cuda()
            model.eval()
            y_hat = model(x)
            test_loss = criterion(y_hat, y)
            running_loss += test_loss
            ps = torch.softmax(y_hat, dim=1)
            predictions = ps.topk(1, dim=1)[1].view(1,-1)[0]
            equals += (predictions == y).sum().item()
            n += len(y)
    test_loss = running_loss/len(testloader)
    accuracy = "Accuracy: {}%\n".format(equals/n*100)
    model.train()
    return (accuracy, test_loss.item())



def train(epoch,  model, criterion, trainloader, testloader, cuda):
    train_loss_list = []
    test_loss_list = []

    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for i in range(epoch):
        running_loss = 0
        for x, y in trainloader:
               if cuda:
                      x,y = x.cuda(), y.cuda()

            optimizer.zero_grad()
            y_hat = model(x)
            print('forward propogated')

            loss = criterion(y_hat, y)
            print('calculated loss')
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        accuracy, test_loss = test(model, testloader, criterion, cuda)
        train_loss = running_loss/len(trainloader)

        train_loss_list.append(train_loss)
        test_loss_list.append(test_loss)

        message = "train-loss: {}    test_loss:{}".format(train_loss, test_loss)
        print(message)
        print(accuracy)
    return (train_loss_list, test_loss_list)


My machine completely stops while training.