external_source doesn't return last incomplete batch

Version

1.31.0

Describe the bug.

An fn.external_source with a callable doesn't seem to return the last batch. I have tried using LastBatchPolicy.PARTIAL, but it seems that the last_batch_policy has no effect on external_source. Is there any way to ensure the last batch is always returned? It seems that this is a bug.

Minimum reproducible example

import numpy as np
from nvidia.dali import pipeline_def
import nvidia.dali.fn as fn
import nvidia.dali.types as types
from nvidia.dali.plugin.pytorch import DALIGenericIterator, LastBatchPolicy

batch_size = 32

class ExternalInputCallable:
    def __init__(self, batch_size):
        self.batch_size = batch_size
        self.length = 100
        self.full_iterations = self.length // batch_size

    def __call__(self, sample_info):
        sample_idx = sample_info.idx_in_epoch
        if sample_idx >= self.length:
            print("STOPPED", sample_idx)
            # Indicate end of the epoch
            raise StopIteration
        return np.array([1], dtype=np.int32), np.array([sample_idx], dtype=np.int32)

@pipeline_def(batch_size=batch_size, num_threads=1, device_id=0)
def callable_pipeline():
    data, label = fn.external_source(source=ExternalInputCallable(batch_size), num_outputs=2, batch=False,
                                       dtype=[types.INT32, types.INT32])
    return data, label

call_pipe = callable_pipeline()
call_pipe.build()

loader = DALIGenericIterator(
    call_pipe,
    last_batch_policy=LastBatchPolicy.PARTIAL,
    output_map=['data', 'label'],
    auto_reset=True,
)
max_label = -1
for batch in loader:
    max_label = max(batch[0]["label"].flatten().max().item(), max_label)
    print(max_label, batch[0]["label"].shape)
print(max_label)

We expect max label to return 99.

Relevant log output

31 torch.Size([32, 1])
STOPPED 100
63 torch.Size([32, 1])
95 torch.Size([32, 1])
95

Other/Misc.

> pip show nvidia-dali-cuda120
Name: nvidia-dali-cuda120
Version: 1.31.0
Summary: NVIDIA DALI  for CUDA 12.0. Git SHA: 166a2e445992d4a0fca41be32ef897ea0a526a6e
Home-page: https://github.com/NVIDIA/dali
Author: NVIDIA Corporation
Author-email: 
License: Apache License 2.0

Check for duplicates

• I have searched the open bugs/issues and have found no duplicates for this bug report