StoreHook.py

# Modified from: https://forums.fast.ai/t/confused-by-output-of-hook-output/29514/4
class StoreHook(HookCallback):
    def on_train_begin(self, **kwargs):
        super().on_train_begin(**kwargs)
        self.hists = []
        
    def hook(self, m, i, o): 
        return o
    
    def on_batch_end(self, train, **kwargs): 
        if (train): 
            self.hists.append(self.hooks.stored[0].cpu().histc(40,0,10))