Improve definition of is_participating

LLNL · Jan 22, 2024 · 2d216ff · 2d216ff
1 parent 96ae91c
commit 2d216ff
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/include/lbann/layers/transform/split.hpp b/include/lbann/layers/transform/split.hpp
@@ -116,7 +116,6 @@ class split_layer : public data_type_layer<TensorDataType>
 
   void fp_setup_outputs() override
   {
-
     const auto& input = this->get_prev_activations();
     auto mini_batch_size =
       this->infer_mini_batch_size_from_parents_or_default_to_current();
@@ -129,6 +128,8 @@ class split_layer : public data_type_layer<TensorDataType>
                                                           El::VC,
                                                           El::ELEMENT,
                                                           Dev> const*>(&input);
+      LBANN_ASSERT(ptr_input);
+
       int tag = 0;
       auto childs = this->get_child_layers();
       if (this->get_communication_flag() == COLL_OPT) {

diff --git a/src/layers/data_type_layer.cpp b/src/layers/data_type_layer.cpp
@@ -87,6 +87,8 @@ data_type_layer<InputTensorDataType, OutputTensorDataType>::operator=(
 template <typename InT, typename OutT>
 bool data_type_layer<InT, OutT>::is_participating() const
 {
+  if (this->subgraph_parallelism_execution())
+    return true;
   if (this->get_num_children() > 0)
     return this->get_activations().Participating();
   if (this->get_num_parents() > 0)