diff --git a/rocAL/rocAL_pybind/amd/rocal/pipeline.py b/rocAL/rocAL_pybind/amd/rocal/pipeline.py
index f30954d639..a53969e61f 100644
--- a/rocAL/rocAL_pybind/amd/rocal/pipeline.py
+++ b/rocAL/rocAL_pybind/amd/rocal/pipeline.py
@@ -217,13 +217,15 @@ def copyToTensorNCHW(self, array,  multiplier, offset, reverse_channels, tensor_
                                        multiplier[0], multiplier[1], multiplier[2], offset[0], offset[1], offset[2], (1 if reverse_channels else 0))
     def GetOneHotEncodedLabels(self, array, device):
         if device=="cpu":
-            return b.getOneHotEncodedLabels(self._handle, ctypes.c_void_p(array.data_ptr()), self._numOfClasses, 0)
+            if (isinstance(array,np.ndarray)):
+                b.getOneHotEncodedLabels(self._handle, array.ctypes.data_as(ctypes.c_void_p), self._numOfClasses, 0)
+            else: #torch tensor
+                return b.getOneHotEncodedLabels(self._handle, ctypes.c_void_p(array.data_ptr()), self._numOfClasses, 0)
         if device=="gpu":
-            return b.getOneHotEncodedLabels(self._handle, ctypes.c_void_p(array.data_ptr()), self._numOfClasses, 1)
-
-    def GetOneHotEncodedLabels_TF(self, array):
-        # Host destination only
-        return b.getOneHotEncodedLabels(self._handle, array.ctypes.data_as(ctypes.c_void_p), self._numOfClasses, 0)
+            if (isinstance(array,cp.ndarray)):
+                b.getOneHotEncodedLabels(self._handle, array.data.ptr, self._numOfClasses, 1)
+            else: #torch tensor
+                return b.getOneHotEncodedLabels(self._handle, ctypes.c_void_p(array.data_ptr()), self._numOfClasses, 1)
 
     def set_outputs(self, *output_list):
         self._output_list_length = len(output_list)
diff --git a/rocAL/rocAL_pybind/amd/rocal/plugin/generic.py b/rocAL/rocAL_pybind/amd/rocal/plugin/generic.py
index ab82b835fc..23d58695e1 100644
--- a/rocAL/rocAL_pybind/amd/rocal/plugin/generic.py
+++ b/rocAL/rocAL_pybind/amd/rocal/plugin/generic.py
@@ -133,7 +133,7 @@ def __next__(self):
         if(self.loader._name == "labelReader"):
             if(self.loader._oneHotEncoding == True):
                 self.loader.GetOneHotEncodedLabels(self.labels, self.device)
-                self.labels_tensor = self.labels.view(-1, self.bs, self.loader._numOfClasses).long()
+                self.labels_tensor = self.labels.reshape(-1, self.bs, self.loader._numOfClasses)
             else:
                 if self.display:
                     for i in range(self.bs):
diff --git a/rocAL/rocAL_pybind/amd/rocal/plugin/tf.py b/rocAL/rocAL_pybind/amd/rocal/plugin/tf.py
index 9f7341fd25..16d95c8ff6 100644
--- a/rocAL/rocAL_pybind/amd/rocal/plugin/tf.py
+++ b/rocAL/rocAL_pybind/amd/rocal/plugin/tf.py
@@ -155,7 +155,7 @@ def __next__(self):
         elif (self.loader._name == "TFRecordReaderClassification"):
             if(self.loader._oneHotEncoding == True):
                 self.labels = np.zeros((self.bs)*(self.loader._numOfClasses),dtype = "int32")
-                self.loader.GetOneHotEncodedLabels_TF(self.labels)
+                self.loader.GetOneHotEncodedLabels(self.labels, device="cpu")
                 self.labels = np.reshape(self.labels, (-1, self.bs, self.loader._numOfClasses))
             else:
                 self.labels = np.zeros((self.bs),dtype = "int32")