A problem about training on GPU

[issyuNaN]

Hi! I am recently trying out the Decipher, it is a really great tool.

However, I noticed that the source code has a little problem when training on other devices.

Here's the detail:

I set the parameter of device as "cuda:0" when ranning the train function, the function compute_v_z_numpy in the path "decipher/tools/_decipher/decipher.py" raised some error. This function accepts an input x, however this tensor x is created on CPU, it did not align its device with the device of paramters of Decipher.

Here I corrected the code as follows:

def compute_v_z_numpy(self, x: np.array):
        
       """Compute decipher_v and decipher_z for a given input.

        Parameters
        ----------
        x : np.ndarray or torch.Tensor
            Input data of shape (n_cells, n_genes).

        Returns
        -------
        v : np.ndarray
            Decipher components v of shape (n_cells, dim_v).
        z : np.ndarray
            Decipher latent z of shape (n_cells, dim_z).
        """
        if type(x) == np.ndarray:
            x = torch.tensor(x, dtype=torch.float32, device=self.device)

        x = torch.log1p(x)
        z_loc, _ = self.encoder_x_to_z(x)
        zx = torch.cat([z_loc, x], dim=-1)
        v_loc, _ = self.encoder_zx_to_v(zx)
        if self.device == "cpu":
            return v_loc.detach().numpy(), z_loc.detach().numpy()
        else:
            return v_loc.cpu().detach().numpy(), z_loc.cpu().detach().numpy()

It can work now. I wonder if you can make some updates or I am allowed to pull a request about this issue?