address comments

pytorch · May 8, 2024 · d0bbf89 · d0bbf89
1 parent d7a6bf4
commit d0bbf89
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_xla/csrc/runtime/cache.h b/torch_xla/csrc/runtime/cache.h
@@ -162,7 +162,6 @@ class PersistentCache : public AbstractCache<K, T, H, E> {
   // if the key is tracked in memory, and if not it will check for a persisted
   // version on disk.
   TypePtr Get(const K& key) override {
-    TORCH_LAZY_TIMED("PersistentCacheGet");
     std::lock_guard<std::mutex> slock(lock_);
     TypePtr mem = memory_cache_.Get(key);
     if (mem) {
@@ -174,6 +173,7 @@ class PersistentCache : public AbstractCache<K, T, H, E> {
       TORCH_LAZY_COUNTER("PersistentCacheMiss", 1);
       return nullptr;
     }
+    TORCH_LAZY_TIMED("PersistentCacheLoad");
     std::stringstream ss;
     std::ifstream in(path, std::ios::binary);
     ss << in.rdbuf();