import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from time import time

# Set device (GPU if available)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")