@@ -3,7 +3,7 @@ from typing import Tuple, List, Mapping, Any, Dict
 
 import numpy as np
 import torch
-from torch.utils.data import default_collate, Dataset
+from torch.utils.data.dataloader import default_collate, Dataset
 
 from super_gradients.common.abstractions.abstract_logger import get_logger
 from super_gradients.common.registry.registry import register_collate_function