Error when running test_fit_model_knn_graph.py on GPU

[anna-lito]

There comes up an error when doing a test run using GPUs. The following error came when pytest tests/test_fit_model_knn_graph.py -s was run from the mubind directory:

==================================== ERRORS ====================================
_________________ ERROR at setup of ModelTests.test_no_errors __________________

cls = <class 'tests.test_fit_model_knn_graph.ModelTests'>
    @classmethod
    def setUpClass(cls):
        warnings.filterwarnings("ignore", category=DeprecationWarning)
        import mubind as mb
        import anndata

        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

        early_stopping = 10
        ad_path = os.path.abspath('tests/_data/pancreas_multiome.h5ad')
        dataloader_path = os.path.abspath('tests/_data/pancreas_multiome.pth')

        ad = anndata.read_h5ad(ad_path)
        cls.train = torch.load(dataloader_path)

        print(ad)
        print(cls.train)

        optimize_log_dynamic = True
        w = [15]
        model = mb.models.Mubind.make_model(cls.train, 4, mb.tl.PoissonLoss(), kernels=[0, 2,] + w, # [0, 2] + w,
                                           # use_dinuc=True, dinuc_mode='full',
                                           optimize_sym_weight=False,
                                           optimize_exp_barrier=True,
                                           optimize_prob_act=True,
                                           optimize_log_dynamic=optimize_log_dynamic,
                                           use_dinuc=False,
                                           device=device,
                                           p_dropout=0.8,
                                           prepare_knn=optimize_log_dynamic,
                                           knn_free_weights=False,
                                           adata=None if not optimize_log_dynamic else ad,
                                           dinuc_mode=None).cuda()

>       cls.model, _ = model.optimize_iterative(cls.train, n_epochs=cls.N_EPOCHS, show_logo=False,
                                                early_stopping=early_stopping, log_each=50,
                                                opt_kernel_shift=0, opt_kernel_length=0,
                                                verbose=1,
                                                use_dinuc=False,
                                                log_next_r2=False)

tests/test_fit_model_knn_graph.py:47: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
mubind/models/models.py:941: in optimize_iterative
    self.optimize_simple(
mubind/models/models.py:675: in optimize_simple
    loss_log_dynamic = self.loss_log_dynamic()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = Mubind(
  (padding): ConstantPad2d(padding=(14, 14, 0, 0), value=0.25)
  (binding_modes): BindingModesSimple(
    (con...taining: [torch.float32 of size 1x250 (GPU 0)]
    )
  )
  (graph_module): GraphModule()
  (criterion): PoissonLoss()
)
    def loss_log_dynamic(self):

        if not hasattr(self.graph_module, 'conn_sparse'):
            return 0

        conn = self.graph_module.conn_sparse

        # log_dynamic = self.graph_module.D_tril # log_dynamic
        # return 100

        # return torch.abs(torch.sparse.sum(self.graph_module.D_tril))

        # log_dynamic = self.graph_module.D_tril.coalesce().values() # self.graph_module.D_tril # log_dynamic
        log_dynamic = self.graph_module.log_dynamic

        idx = conn.indices()
        conn_vals = conn.values()
        pos = torch.arange(idx.size(1))

        # prepare combinations based on common indexes
        uniq_idx = idx.unique()
        all_combinations = []
        for u_idx in uniq_idx:
            # at least one common index has to be present in the position retrieved
>           sub_pos = pos[(idx[0] == u_idx) | (idx[1] == u_idx)]
E           RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

mubind/models/models.py:441: RuntimeError
=========================== short test summary info ============================
ERROR tests/test_fit_model_knn_graph.py::ModelTests::test_no_errors - Runtime...
============================== 1 error in 42.17s ===============================

[ilibarra]

Hi @anna-lito!

I tested the updated main branch, and this error is not shown anymore while executing pytest. Please feel free to test again and close, if no more local errors. Once pytests are tested by GitHub, then I'll close this issue.

Thanks!

[ilibarra]

Solved in a232a9f. Feel free to reopen if issues.