diff --git a/vitookit/datasets/ffcv_transform.py b/vitookit/datasets/ffcv_transform.py
index a8ea5890d7a92f5d25b1efe78e13a481d80c98c1..036b5933a48720ea48be09229cb40f9f801a6363 100644
--- a/vitookit/datasets/ffcv_transform.py
+++ b/vitookit/datasets/ffcv_transform.py
@@ -356,7 +356,7 @@ def convolution(image: np.ndarray, kernel: list | tuple, output: np.ndarray) ->
     else:
         raise Exception('Shape of image not supported')
 
-    m_k, n_k, _ = kernel.shape
+    m_k, n_k = kernel.shape
 
     y_strides = m_i - m_k + 1  # possible number of strides in y direction
     x_strides = n_i - n_k + 1  # possible number of strides in x direction
diff --git a/vitookit/evaluation/eval_cls_ffcv.py b/vitookit/evaluation/eval_cls_ffcv.py
index b71e3aa99207fba4aa80c0868d3e58c27ae10bd5..70a79f0645280dda6fffc0507d2f83c60317ed7b 100644
--- a/vitookit/evaluation/eval_cls_ffcv.py
+++ b/vitookit/evaluation/eval_cls_ffcv.py
@@ -263,7 +263,7 @@ def main(args):
     misc.init_distributed_mode(args)
 
     print(args)
-
+    import torch
     device = torch.device(args.device)
 
     # fix the seed for reproducibility
diff --git a/vitookit/evaluation/eval_linear_ffcv.py b/vitookit/evaluation/eval_linear_ffcv.py
index 26cabe0db90025c720c8f88d23efc12feecc9e9b..143b6a8f7fa58120aa56fad70d9265dbdeb781bf 100644
--- a/vitookit/evaluation/eval_linear_ffcv.py
+++ b/vitookit/evaluation/eval_linear_ffcv.py
@@ -126,7 +126,7 @@ def main(args):
     print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
     print("{}".format(args).replace(', ', ',\n'))
 
-    import torch
+    
     device = torch.device(args.device)
 
     # fix the seed for reproducibility