update __call__ method in rbflayer

Eloypripan · Dec 9, 2020 · be2a74d · be2a74d
1 parent 2410dca
commit be2a74d
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/rbflayer.py b/rbflayer.py
@@ -17,9 +17,14 @@ def __init__(self, X):
         self.X = X
 
     def __call__(self, shape, dtype=None):
-        assert shape[1] == self.X.shape[1]
-        idx = np.random.randint(self.X.shape[0], size=shape[0])
-        return self.X[idx, :]
+	assert shape[1] == self.X.shape[1]
+  	idx = np.random.randint(self.X.shape[0], size=shape[0])
+
+	# type checking to access elements of data correctly
+  	if type(self.X) == np.ndarray:
+    		return self.X[idx, :]
+  	elif type(self.X) == pd.core.frame.DataFrame:
+    		return self.X.iloc[idx, :]
 
 
 class RBFLayer(Layer):