[TORCH] FED DEBUG

leondavi · Aug 3, 2024 · c9dfa6a · c9dfa6a
1 parent 5a1038e
commit c9dfa6a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src_erl/NerlnetApp/src/Bridge/onnWorkers/workerGeneric.erl b/src_erl/NerlnetApp/src/Bridge/onnWorkers/workerGeneric.erl
@@ -208,6 +208,13 @@ wait(cast, {end_stream , StreamName}, State = #workerGeneric_state{myName = _MyN
 
 wait(cast, {post_train_update, Data}, State = #workerGeneric_state{myName = _MyName, distributedBehaviorFunc = DistributedBehaviorFunc}) ->
   NextStateBehavior = DistributedBehaviorFunc(post_train, {get(generic_worker_ets), {post_train_update, Data}}),
+  if 
+    NextStateBehavior == train -> 
+      ok;
+    true -> 
+      ?LOG_ERROR("@wait: post_train controller method must return train atom!"),
+      throw("@wait: post_train controller method must return train atom!")
+  end,
   handle_end_stream_waiting_list(DistributedBehaviorFunc, train),
   {next_state, NextStateBehavior, State};