small fixes

source/NN/model.py

@@ -1,4 +1,6 @@
 import torch.nn as nn
+import torch
+import torch.nn.functional as F
 
 
 class Neural_Network_Model(nn.Module):
@@ -16,5 +18,4 @@ class Neural_Network_Model(nn.Module):
         x = self.fc2(x)
         x = torch.relu(x)
         x = self.out(x)
-        F.log_softmax(x, dim=-1)
-        return x
+        return F.log_softmax(x, dim=-1)

source/NN/neural_network.py

@@ -1,15 +1,16 @@
 import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
-from torchvision import datasets, transforms
+from torchvision import datasets, transforms, utils
 from torchvision.transforms import Compose, Lambda, ToTensor
 import matplotlib.pyplot as plt
 import numpy as np
 from model import *
 
+device = torch.device('cuda')
+
 #data transform to tensors:
-data_transformer = transforms.Compose
-([
+data_transformer = transforms.Compose([
     transforms.Resize((150, 150)),  
     transforms.ToTensor(),  
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  
@@ -46,6 +47,9 @@ def train(model, dataset, iter=100, batch_size=64):
             loss = criterion(output, labels.to(device))
             loss.backward()
             optimizer.step()
+        if epoch % 10 == 0:
+                print('epoch: %3d loss: %.4f' % (epoch, loss))
+
 #function for getting accuracy
 def accuracy(model, dataset):
     model.eval()
@@ -57,6 +61,9 @@ def accuracy(model, dataset):
     return correct.float() / len(dataset)
 
 
+
+
 model = Neural_Network_Model()
+model.to(device)
 train(model, train_set)
 print(accuracy(model, test_set))