version change

kchu25 · Jun 19, 2022 · 3b06afa · 3b06afa
1 parent e35fcdc
commit 3b06afa
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "FastaLoader"
 uuid = "139838d8-4077-4d8a-94e1-e6dd554a184c"
 authors = ["Shane Kuei Hsien Chu (skchu@wustl.edu)"]
-version = "0.1.6"
+version = "0.1.7"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

diff --git a/src/fasta_w_splits.jl b/src/fasta_w_splits.jl
@@ -147,11 +147,21 @@ end
 
 function get_test_set_for_flux(fws::FASTA_DNA_w_splits; gpu=true)
     test_set_ind = get_test_set_ind(fws.mcs);
-    return fws.data_matrix_gpu[:,:,test_set_ind], fws.label_indicators_gpu[test_set_ind,:]
+    if gpu
+        return fws.data_matrix_gpu[:,:,test_set_ind], fws.label_indicators_gpu[test_set_ind,:]
+    else
+        return fws.data_matrix[:,:,test_set_ind], fws.label_indicators[test_set_ind,:]
+    end
 end
 
 function get_train_fold_for_flux(fws::FASTA_DNA_w_splits, fold::Int; gpu=true)
     train_set_ind, valid_set_ind = get_train_fold_ind(fws.mcs, fold)
-    return fws.data_matrix_gpu[:,:,train_set_ind], fws.label_indicators_gpu[valid_set_ind,:]
+    if gpu
+        return fws.data_matrix_gpu[:,:,train_set_ind],  fws.label_indicators_gpu[train_set_ind,:]
+               fws.data_matrix_gpu[:,:,valid_set_ind],  fws.label_indicators_gpu[valid_set_ind,:]
+    else
+        return fws.data_matrix[:,:,train_set_ind], fws.label_indicators[train_set_ind,:]
+               fws.data_matrix[:,:,valid_set_ind], fws.label_indicators[valid_set_ind,:]
+    end
 end