import time, torch, torchvision

DEVICE = "cuda"
model = torchvision.models.resnet18().to(DEVICE).train()
optimizer = torch.optim.Adam(model.parameters())