partial.py

#check gpu availability, can force cpu use with -1 as gpu                                                                                                                                                       
    cuda = torch.cuda.is_available()
    if cuda and gpu_to_use > -1:
        use_cuda = True
        with torch.cuda.device(gpu_to_use):
            torch.cuda.seed()
            device = torch.device("cuda", gpu_to_use)
        print("Using CUDA.")
    else:
        use_cuda = False
        device = torch.device("cpu")
        print("Using CPU.")

    net = model.to(device)

    try:
        save_file_loss = save_file + "_loss"
        #training the model going through the data several times                                                                                                                                                    
        for i in range(epochs):
            #get the loss on the epoch                                                                                                                                                                              
            print ("Current Epoch:{}".format(i))
            if model_type != 5 and model_type != 6:
                training_res = de.train(net, dataloader, BATCH_SIZE, MAX_SIZE, device, learning_rate, momentum, optimizer)

                if args.test:
                    testing_res = de.test(net, test_dataloader, test_batch, MAX_SIZE, device)
                else:
                    testing_res = []
                ut.save_data(save_file_loss, training_res, testing_res, i)
            else:
                training_res = de.train_rnn(net, dataloader, BATCH_SIZE, device, learning_rate, hidden)
                #TODO: testing for rnn                                                                                                                                                                              
                ut.save_data(save_file_loss, training_res, [], i)