fix a bug when specifying batch size in RModel with input param bs

Author: lmoneta

tmva/sofie/src/RModel.cxx

@@ -263,23 +263,37 @@ void RModel::Initialize(int batchSize, bool verbose) {
    // if the batch size is provided it can be used to specify the full shape
    // Add the full specified tensors in fReadyInputTensors collection
    for (auto &input : fInputTensorInfos) {
+      std::cout << "looking at the tensor " << input.first << std::endl;
       // if a batch size is provided convert batch size
       // assume is parametrized as "bs" or "batch_size"
       if (batchSize > 0) {
          // std::vector<Dim> shape;
          // shape.reserve(input.second.shape.size());
-         for (auto &d : input.second.shape) {
-            if (d.isParam && (d.param == "bs" || d.param == "batch_size")) {
-               d = Dim{static_cast<size_t>(batchSize)};
+         // assume first parameter is teh batch size
+         if (!input.second.shape.empty()) {
+            auto & d0 = input.second.shape[0];
+            if (d0.isParam) {
+               if (verbose) std::cout << "Fix the batch size to " << batchSize << std::endl;
+               d0 = Dim{static_cast<size_t>(batchSize)};
+            }
+            else {  // look for cases that a bs or bath_size is specified in tensor shape
+               for (auto &d : input.second.shape) {
+                  if (d.isParam && (d.param == "bs" || d.param == "batch_size")) {
+                     d = Dim{static_cast<size_t>(batchSize)};
+                     if (verbose) std::cout << "Input shape has bs or batch_size as names. Fix the batch size to " << batchSize << std::endl;
+                  }
+               }
             }
          }
       }
       auto shape = ConvertShapeToInt(input.second.shape);
       if (!shape.empty()) {
-         // add to the ready input tensor informations
-         AddInputTensorInfo(input.first, input.second.type, shape);
-         // remove from the tensor info
+#if 0
+         // remove from the tensor info old dynamic shape
          fInputTensorInfos.erase(input.first);
+         // add to the ready input tensor information the new fixed shape
+         AddInputTensorInfo(input.first, input.second.type, shape);  
+#endif
       }
       // store the parameters of the input tensors
       else {