Commit bf86ea9

committed

[Fix] device error

1 parent eddffdf commit bf86ea9Copy full SHA for bf86ea9

File tree

-2

lines changed

-2

lines changed

Lines changed: 1 addition & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -82,7 +82,7 @@ def __init__(self, dims: List[int], device: str) -> None:`
`82`	`82`	`) for i in range(len(dims)-1)`
`83`	`83`	`]`
`84`	`84`
`85`		`- self.classify_layer = torch.nn.Linear(dims[-1], 10)`
	`85`	`+ self.classify_layer = torch.nn.Linear(dims[-1], 10).to(device)`
`86`	`86`	`self.criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.1)`
`87`	`87`	`self.optimizer = torch.optim.AdamW(self.classify_layer.parameters(), lr=0.01)`
`88`	`88`	`self.softmax = torch.nn.Softmax(dim=1)`

Lines changed: 1 addition & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -36,9 +36,9 @@ def main() -> None:`
`36`	`36`	`loss_logger.reset()`
`37`	`37`
`38`	`38`	`for inputs, labels in pbar:`
`39`		`- break`
`40`	`39`	`pos_inputs = create_pos_data(inputs, labels).to(device)`
`41`	`40`	`neg_inputs = create_neg_data(inputs, labels).to(device)`
	`41`	`+ labels = labels.to(device)`
`42`	`42`
`43`	`43`	`loss = model(pos_inputs=pos_inputs, neg_inputs=neg_inputs, pos_labels=labels)`
`44`	`44`	`loss_logger.update(loss, inputs.shape[0])`

Comments

(0)