model.to(device)
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)