training.py

import torch
from headshot import Headshot
from headshot import config


num_epochs = 30
out_dir = 'out'
device = 'cuda' if torch.cuda.is_available() else 'cpu'





model = Headshot().to(device)

model.train()