From 7215b4ec5b0193e88bf431b47c14824026563d25 Mon Sep 17 00:00:00 2001
From: meyerkm <129152803+meyerkm@users.noreply.github.com>
Date: Tue, 26 Sep 2023 09:37:08 +0200
Subject: [PATCH] Update train.py DataLoader (#25)

Bug fix groupby action for categorical data to only show observed values in resulting DataFrame.
---
 deeprvat/deeprvat/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeprvat/deeprvat/train.py b/deeprvat/deeprvat/train.py
index fd047868..c2e2abd4 100644
--- a/deeprvat/deeprvat/train.py
+++ b/deeprvat/deeprvat/train.py
@@ -285,7 +285,7 @@ def __getitem__(self, index):
         start_idx = index * self.batch_size
         end_idx = min(self.total_samples, start_idx + self.batch_size)
         batch_samples = self.sample_order.iloc[start_idx:end_idx]
-        samples_by_pheno = batch_samples.groupby("phenotype")
+        samples_by_pheno = batch_samples.groupby("phenotype", observed=True)
 
         result = dict()
         for pheno, df in samples_by_pheno: