Update 'data_sampler.py'

data_sampler.py

     """
     Samples train/test data and masks based on row indices.
     """
-    def __init__(
-        self,
-        original_adj_mat: np.ndarray,
-        null_mask: np.ndarray,
-        train_index: np.ndarray,
-        test_index: np.ndarray
-    ) -> None:
+    def __init__(self, original_adj_mat, null_mask, train_index, test_index):
+        super(ExterSampler, self).__init__()
         self.adj_mat = original_adj_mat
         self.null_mask = null_mask
         self.train_index = train_index
         self.test_index = test_index
-        self.train_data, self.test_data = self._sample_train_test_data()
-        self.train_mask, self.test_mask = self._sample_train_test_mask()
+        self.train_data, self.test_data = self.sample_train_test_data()
+        self.train_mask, self.test_mask = self.sample_train_test_mask()
 
-    def _sample_train_test_data(self) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Samples train and test data based on row indices."""
-        test_data = self.adj_mat.copy()
+    def sample_train_test_data(self):
+        # If self.adj_mat is already a Tensor, convert it to a numpy array
+        if isinstance(self.adj_mat, torch.Tensor):
+            adj = self.adj_mat.cpu().numpy().copy()
+        else:
+            adj = self.adj_mat.copy()
+
+        # Create test_data by copying and modifying the rows corresponding to train_index
+        test_data = adj.copy()
         test_data[self.train_index, :] = 0
-        train_data = self.adj_mat - test_data
-        return torch.from_numpy(train_data), torch.from_numpy(test_data)
+        train_data = adj - test_data
 
-    def _sample_train_test_mask(self) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Creates train and test masks with negative sampling."""
-        neg_value = np.ones(self.adj_mat.shape, dtype=np.float32) - self.adj_mat - self.null_mask
-        neg_train = neg_value.copy()
+        # Convert back to Tensor (since subsequent operations are performed on Tensors)
+        train_data = torch.from_numpy(train_data)
+        test_data = torch.from_numpy(test_data)
+        return train_data, test_data
+
+    def sample_train_test_mask(self):
+        # Ensure that adj_mat and null_mask are numpy arrays
+        if isinstance(self.adj_mat, torch.Tensor):
+            adj = self.adj_mat.cpu().numpy()
+        else:
+            adj = self.adj_mat
+        if isinstance(self.null_mask, torch.Tensor):
+            null_mask = self.null_mask.cpu().numpy()
+        else:
+            null_mask = self.null_mask
+
+        neg_value = np.ones(adj.shape, dtype=np.float32)
+        neg_train = neg_value - adj - null_mask
         neg_train[self.test_index, :] = 0
-        neg_test = neg_value.copy()
+        neg_test = neg_value - adj - null_mask
         neg_test[self.train_index, :] = 0
 
-        train_mask = (self.train_data.numpy() + neg_train).astype(bool)
-        test_mask = (self.test_data.numpy() + neg_test).astype(bool)
-        return torch.from_numpy(train_mask), torch.from_numpy(test_mask)
+        train_data_np = self.train_data.cpu().numpy() if isinstance(self.train_data, torch.Tensor) else self.train_data
+        test_data_np = self.test_data.cpu().numpy() if isinstance(self.test_data, torch.Tensor) else self.test_data
+
+        train_mask = (train_data_np + neg_train).astype(bool)
+        test_mask = (test_data_np + neg_test).astype(bool)
+        train_mask = torch.from_numpy(train_mask)
+        test_mask = torch.from_numpy(test_mask)
+        return train_mask, test_mask
+    
     
     
 class RegressionSampler(object):