BugFix - LBP (#216)

Co-authored-by: Swetha B S <swetha@multicorrewareinc.com>
Co-authored-by: Kiriti Gowda <kiritigowda@gmail.com>
Co-authored-by: Lakshmi Kumar <lakshmi.kumar@amd.com>

rocAL/source/readers/file_source_reader.cpp

@@ -42,31 +42,21 @@ FileSourceReader::FileSourceReader() {
 
 unsigned FileSourceReader::count_items() {
     int ret = 0;
+    int size = 0;
     if (_shard_size == -1) {                                     // When shard_size is set to -1, The shard_size variable is not used
-        if (_loop) return largest_shard_size_without_padding();  // Return the size of the largest shard amongst all the shard's size
+        if (_loop) return largest_shard_size_without_padding();  // Return the size of the largest shard amongst all the shard's size - 
+                                                                 // _file_count_all_shards is now padded - hence with/ without padding is one and the same
         int size = std::max(largest_shard_size_without_padding(), _batch_size);
-        ret = (size - _read_counter);
-        // Formula used to calculate - [_last_batch_padded_size = _batch_size - (_shard_size % _batch_size) ]
-        // Since "size" doesnt involve padding - we add the count of padded samples to the number of remaining elements
-        // which equals to the shard size with padding
-        if (_last_batch_info.last_batch_policy == RocalBatchPolicy::PARTIAL || _last_batch_info.last_batch_policy == RocalBatchPolicy::FILL) {
-            ret += _last_batch_padded_size;
-        } else if (_last_batch_info.last_batch_policy == RocalBatchPolicy::DROP &&
-                   _last_batch_info.pad_last_batch_repeated == true) {  // When pad_last_batch_repeated is False - Enough
-                                                       // number of samples would not be present in the last batch - hence
-                                                       // dropped by condition handled in the loader
-            ret -= _batch_size;
-        }
     } else if (_shard_size > 0) {
-        auto largest_shard_size_with_padding =
+        auto largest_shard_size_with_padding = 
             _shard_size + (_batch_size - (_shard_size % _batch_size));  // The shard size used here is padded
-        if (_loop)
+        if (_loop) 
             return largest_shard_size_with_padding;
         int size = std::max(largest_shard_size_with_padding, _batch_size);
-        ret = (size - _read_counter);
-        if (_last_batch_info.last_batch_policy == RocalBatchPolicy::DROP)  // The shard size is padded at the beginning of the condition, hence dropping the last batch
-            ret -= _batch_size;
     }
+    ret = (size - _read_counter);
+    if (_last_batch_info.last_batch_policy == RocalBatchPolicy::DROP && _last_batch_padded_size != 0)
+        ret -= _batch_size;
     return ((ret < 0) ? 0 : ret);
 }
 
@@ -272,7 +262,7 @@ Reader::Status FileSourceReader::generate_file_names() {
     auto dataset_size = _file_count_all_shards;
     // Pad the _file_names with last element of the shard in the vector when _pad_last_batch_repeated is True
     _padded_samples = ((_shard_size > 0) ? _shard_size : largest_shard_size_without_padding()) % _batch_size;
-    _last_batch_padded_size = (_batch_size > 1) ? (_batch_size - _padded_samples) : 0;
+    _last_batch_padded_size = ((_batch_size > 1) && (_padded_samples > 0 )) ? (_batch_size - _padded_samples) : 0;
 
     if (_pad_last_batch_repeated == true) {
                                             // pad the last sample when the dataset_size is not divisible by
@@ -366,7 +356,6 @@ void FileSourceReader::compute_start_and_end_idx_of_all_shards() {
         auto end_idx_of_shard = start_idx_of_shard + actual_shard_size_without_padding() - 1;
         _shard_start_idx_vector.push_back(start_idx_of_shard);
         _shard_end_idx_vector.push_back(end_idx_of_shard);
-
     }
 }