data_x, data_y = train_x, train_y
elif (args.dataset == "test"):
data_x, data_y = test_x, test_y
elif (args.dataset == "valid"):
data_x, data_y = valid_x, valid_y
test_dataset = Dataset(root_dir, data_x, data_y, transforms=transform)
test_generator = torch.utils.data.DataLoader(test_dataset, **params)
print("Loaded dataloaders...")
criterion = torch.nn.CrossEntropyLoss()
model = NeuralNet(0.001, criterion, 64, 2)
state_dict = torch.load(model_name)
model.load_state_dict(state_dict)
for parameter in model.parameters():
parameter.requires_grad = False
if (use_cuda):
model.cuda()
model.eval()
summary(model, (1, 64, 64))
print("Loaded model...")
preds = []
labels = []
for local_batch, local_labels in tqdm(test_generator):
labels.extend(local_labels.numpy().tolist())
local_batch, local_labels = local_batch.to(device), local_labels.to(device)
output = model.forward(local_batch)
