logo

classifier-build-train

# TODO: Build and train your network

# check for gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# network instance
neural_net = net_from_torchvision([1024,512], 102, 'relu', device, learn_rate = 0.001)

# train for 25 epochs
neural_net.train_network(dataloaders['train'], dataloaders['validate'], 5, plot = True)
  • Share