Better support for cuda devices

src/utils.py

@@ -14,7 +14,6 @@ import pandas as pd
 import torch
 import torch.nn as nn
 import yaml
-from torch import mode
 from torch.optim import Optimizer
 
 from src.batch_loading import get_batches, get_ordered_dataframe_len