Update train_model.py

Tobias-Fischer · Aug 6, 2020 · 900c78b · 900c78b
1 parent ece1e8d
commit 900c78b
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/rt_gene_model_training/pytorch/train_model.py b/rt_gene_model_training/pytorch/train_model.py
@@ -173,13 +173,13 @@ def test_dataloader(self):
             _train_subjects.append([1, 2, 8, 10, 5, 6, 11, 12, 13])
             _train_subjects.append([3, 4, 7, 9, 5, 6, 11, 12, 13])
             # validation set is always subjects 14, 15 and 16
-            _valid_subjects.append([14, 15, 16])
-            _valid_subjects.append([14, 15, 16])
-            _valid_subjects.append([14, 15, 16])
+            _valid_subjects.append([0, 14, 15, 16])
+            _valid_subjects.append([0, 14, 15, 16])
+            _valid_subjects.append([0, 14, 15, 16])
             # test subjects
-            _test_subjects.append([0, 5, 6, 11, 12, 13])
-            _test_subjects.append([0, 3, 4, 7, 9])
-            _test_subjects.append([0, 1, 2, 8, 10])
+            _test_subjects.append([5, 6, 11, 12, 13])
+            _test_subjects.append([3, 4, 7, 9])
+            _test_subjects.append([1, 2, 8, 10])
         else:
             _train_subjects.append([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
             _valid_subjects.append([0])  # Note that this is a hack and should not be used to get results for papers