diff --git a/torch_modules/src/ConfigDataset.cpp b/torch_modules/src/ConfigDataset.cpp index 3a4e6466710340d24d5239bb9105cb4f3b2f4222..30dce0e13296b25ea4b4391d7abd2926c327136d 100644 --- a/torch_modules/src/ConfigDataset.cpp +++ b/torch_modules/src/ConfigDataset.cpp @@ -137,6 +137,10 @@ void ConfigDataset::computeNbToGive() smallestSize = sizeLeft; } for (auto & it : holders) - nbToGive[it.first] = std::floor(1.0*it.second.sizeLeft()/smallestSize); + { + nbToGive[it.first] = std::max<int>(1,std::floor(1.0*it.second.sizeLeft()/smallestSize)); + if (it.second.sizeLeft() == 0) + nbToGive[it.first] = 0; + } }