Fix bug in caching allocator. (#43719)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43719
Accidentally this slipped through: with guard did not update the current
context
Test Plan: cpu_caching_allocator_test
Reviewed By: linbinyu
Differential Revision: D23374453
fbshipit-source-id: 1d3ef21cc390d0a8bde98fb1b5c2175b40ab571b
