This is my entire code for a small classification problem,The problem is that when i run this it consumes a lot of CPU ,I am not sure which part of the code is responsible for it,
Any ideas or suggestions will be really helpful,
Thanks in advance ,
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import torch.optim as optim
import torchvision
from torch.optim import lr_scheduler
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
from tensorboardX import SummaryWriter
import setproctitle
from tensorboard_logger import log_value, configure

RUN_PATH = '/mnt/da5df9e4-cdc6-4d55-91e8-b2383e89165f/Ryan/1234/models'
configure(RUN_PATH, flush_secs=1)
setproctitle.setproctitle('train.py')
writer =  SummaryWriter('/mnt/machine/Ryan/1234/models')
    
net = Net().cuda()
print (net)
transform = transforms.Compose([
    transforms.Scale((224,224)),
    transforms.ToTensor()
])

class WhaleData(Dataset):
    def __init__(self, data_file, root_dir , transform = None):
        self.csv_file = pd.read_csv(data_file)
        self.root_dir = root_dir
        self.transform = transform
        
    def __len__(self):
        return len(os.listdir(self.root_dir))
    
    def __getitem__(self, index):
        image = os.path.join(self.root_dir, self.csv_file['Image'][index])
        image = Image.open(image).convert('RGB')
        label  = self.csv_file['Id'][index]
        sample = {'image': image, 'label':label}
        if self.transform:
            sample['image'] = self.transform(sample['image'])
        return sample

out_path = '/mnt/machine/Ryan/1234/models/'
trainset  = WhaleData(data_file = '/mnt/machine/Ryan/1234/train_encoded.csv', 
         root_dir = '/mnt/machine/Ryan/1234/train',transform = transform)
train_loader = torch.utils.data.DataLoader(trainset,batch_size = 128)
test_loader = torch.utils.data.DataLoader(trainset)

model_ft = torchvision.models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 4251)
model_ft = model_ft
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

# criterion = nn.CrossEntropyLoss()
# optimizer = optim.SGD(net.parameters(), lr = 0.001, momentum = 0.9)
for epoch in range(100):
    running_loss = 0.0
    for index,info in enumerate(test_loader):
        inputs = info['image']
        inputs = Variable(inputs.float())#.cuda()
        labels= Variable(info['label'].long())#.cuda()
        # print (model_conv)
        outputs = model_ft(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss+=loss
        print ('Loss for this image  ', loss.data[0])
        writer.add_scalar('loss', loss.data[0])
        log_value('train_loss', loss, epoch)
    torch.save(model_ft.state_dict, out_path+'{}.pth'.format(epoch))