Implement better caching allocator for segmentation usecase. (#44618)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/44618
This diff refactors caching allocator to allow for overriding behavior by
making it a virtual class.
Test Plan: https://www.internalfb.com/intern/fblearner/details/218419618?tab=Experiment%20Results
Reviewed By: dreiss
Differential Revision: D23672902
fbshipit-source-id: 976f02922178695fab1c87f453fcb59142c258ec
