Resolve malaria dataset error (#86)

amzn · Apr 19, 2022 · aed9a89 · aed9a89
1 parent dd4a6a2
commit aed9a89
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/finite_ntk/data.py b/finite_ntk/data.py
@@ -52,7 +52,7 @@ def generate_data(
 
     df = pd.read_hdf(hdf_loc, "full")
 
-    is_train_year = torch.from_numpy((df["year"] == test_year).values)
+    is_train_year = torch.from_numpy((df["year"] == train_year).values)
     is_ng = torch.from_numpy(df["is_ng"].values).bool()
     is_test = torch.from_numpy((df["year"] == test_year).values)