Merge branch 'adding_da_to_models' of github.com:USGS-R/river-dl into…

… adding_da_to_models
USGS-R · Jun 4, 2021 · e939460 · e939460
2 parents 3e54465 + ab4cff1
commit e939460
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/river_dl/loss_functions.py b/river_dl/loss_functions.py
@@ -80,12 +80,8 @@ def nnse_one_var_samplewise(data, y_pred, var_idx, tasks):
 
 
 def y_data_components(data, y_pred, var_idx, tasks):
-    if tasks == 2: 
-        weights = data[:, :, -2:]
-        y_true = data[:, :, :-2]
-    else: 
-        weights = data[:, :, -1:]
-        y_true = data[:, :, :-1]
+    weights = data[:, :, -tasks:]
+    y_true = data[:, :, :-tasks]
 
     # ensure y_pred, weights, and y_true are all tensors the same data type
     y_true = tf.convert_to_tensor(y_true)