diff --git a/TrainingExtensions/torch/src/python/aimet_torch/custom/custom_tensor_utils.py b/TrainingExtensions/torch/src/python/aimet_torch/custom/custom_tensor_utils.py
index 29d13d4e03..2682d0329d 100644
--- a/TrainingExtensions/torch/src/python/aimet_torch/custom/custom_tensor_utils.py
+++ b/TrainingExtensions/torch/src/python/aimet_torch/custom/custom_tensor_utils.py
@@ -45,7 +45,7 @@ def to_torch_tensor(tensors):
 
     def to_custom_tensor(original, torch_tensors):
         """ placeholder in case spconv doesn't exist """
-        return original
+        return torch_tensors
 else:
     from typing import List, Union, Tuple
     import torch