Hook_output seems to be missing last batch of training data

Part 1 (2019) Advanced (Part 1 v3)
1

Running this on 1.0.46 …

class StoreHook(HookCallback):
    def on_train_begin(self, **kwargs):
        super().on_train_begin(**kwargs)
        self.train_outputs = []
        self.val_outputs = []
        
    def hook(self, m, i, o): 
        return o
    
    def on_batch_end(self, train, **kwargs): 
        if (train): 
            self.train_outputs.append(self.hooks.stored[0])
        else:
            self.val_outputs.append(self.hooks.stored[0])

cb = learn.callbacks[0]
torch.cat(cb.train_outputs).shape, torch.cat(cb.val_outputs).shape
# => (torch.Size([1792, 512]), torch.Size([455, 512]))
len(data.train_ds), len(data.valid_ds), len(data.train_dl), len(data.valid_dl)
# => (1822, 455, 28, 8)

I’m expected the callback to return 1822 activations (the same number of items in the training dataset) but it is only returning 1792. I’ve added the obligatory pdb.set_trace() and print statements in the on_batch_end method and I can verify that the final 30 training images (the last batch) is never passed to the method.

Any ideas?

This was working way back in the day so perhaps something has changed with the API?

Thanks - wg