diff --git a/pytorch_toolbelt/utils/torch_utils.py b/pytorch_toolbelt/utils/torch_utils.py
index 478fe4a77..61ae84244 100644
--- a/pytorch_toolbelt/utils/torch_utils.py
+++ b/pytorch_toolbelt/utils/torch_utils.py
@@ -47,6 +47,7 @@
     "get_collate_for_dataset",
     "get_non_wrapped_model",
     "container_to_tensor",
+    "convert_2d_to_3d",
 ]