Commit 305a910

committed

[Fix] device error

1 parent bf86ea9 commit 305a910Copy full SHA for 305a910

File tree

-3

lines changed

-3

lines changed

Lines changed: 3 additions & 3 deletions

Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ def main() -> None:`
`17`	`17`
`18`	`18`	`# DataLoader`
`19`	`19`	`train_dataloader = get_mnist_dataloader(_mode="train", batch_size=batch_size)`
`20`		`- val_dataloader = get_mnist_dataloader(_mode="val", batch_size=1)`
	`20`	`+ val_dataloader = get_mnist_dataloader(_mode="val", batch_size=16)`
`21`	`21`
`22`	`22`	`# Device`
`23`	`23`	`device = "cuda:0" if torch.cuda.is_available() else "cpu"`
`@@ -53,9 +53,9 @@ def main() -> None:`
`53`	`53`	`predicts = []`
`54`	`54`	`targets = []`
`55`	`55`	`for inputs, labels in tqdm(val_dataloader):`
`56`		`- inputs_all_labels = create_test_data(inputs).to(device)`
	`56`	`+ # inputs_all_labels = create_test_data(inputs).to(device)`
`57`	`57`
`58`		`- predict = model.predict(inputs_all_labels)`
	`58`	`+ predict = model.predict(inputs)`
`59`	`59`
`60`	`60`	`predicts.extend(predict.tolist())`
`61`	`61`	`targets.append(labels.item())`

Comments

(0)