[FIX] modified for checking weight grad

This path checks requested memory is weight gradient which information
will be used for planning.

Signed-off-by: Jiho Chu <jiho.chu@samsung.com>

diff --git a/nntrainer/tensor/cache_pool.cpp b/nntrainer/tensor/cache_pool.cpp
index ac14d99..3a4c188 100644 (file)
--- a/nntrainer/tensor/cache_pool.cpp
+++ b/nntrainer/tensor/cache_pool.cpp
@@ -154,9 +154,9 @@ void CachePool::invalidate(unsigned int id) {
 unsigned int CachePool::requestMemory(size_t bytes, unsigned int start_time,
                                       unsigned int end_time,
                                       std::vector<unsigned int> exec_order,
-                                      TensorLifespan lifespan) {
+                                      TensorLifespan lifespan, bool is_wgrad) {
   auto id = MemoryPool::requestMemory(bytes, start_time, end_time, exec_order,
-                                      lifespan);
+                                      lifespan, is_wgrad);
 
   const CachePolicy policy = convertTensorLifespanToCachePolicy(lifespan);
 

diff --git a/nntrainer/tensor/cache_pool.h b/nntrainer/tensor/cache_pool.h
index b78fa00..8514ebe 100644 (file)
--- a/nntrainer/tensor/cache_pool.h
+++ b/nntrainer/tensor/cache_pool.h
@@ -75,7 +75,8 @@ public:
   virtual unsigned int requestMemory(
     size_t bytes, unsigned int start_time, unsigned int end_time,
     std::vector<unsigned int> exec_order = std::vector<unsigned int>(),
-    TensorLifespan lifespan = TensorLifespan::MAX_LIFESPAN);
+    TensorLifespan lifespan = TensorLifespan::MAX_LIFESPAN,
+    bool is_wgrad = false);
 
   /**
    * @brief Get the allocated cache