From 5260044220ed61dd9a307cda51911209067927a1 Mon Sep 17 00:00:00 2001
From: Nirmayi <ynirmayi@gmail.com>
Date: Wed, 20 Nov 2024 17:40:47 +0100
Subject: [PATCH] attempt to update docker setup for gpu support

---
 src/methods/destvi/config.vsh.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/methods/destvi/config.vsh.yaml b/src/methods/destvi/config.vsh.yaml
index c81288f..7d2f87f 100644
--- a/src/methods/destvi/config.vsh.yaml
+++ b/src/methods/destvi/config.vsh.yaml
@@ -35,6 +35,7 @@ engines:
     setup:
       - type: python
         packages: 
+          - jax[cuda12]
           - scvi-tools
   
 runners: