Fix memory leak

Author: Critsium-xy

source/module_base/memory.cpp

@@ -23,7 +23,7 @@ int Memory::short_memory = sizeof(short); // 2.0 Byte
 
 int Memory::n_memory = 1000;
 int Memory::n_now = 0;
-bool Memory::init_flag =  false;
+bool Memory::init_flag = false;
 
 #if defined(__CUDA) || defined(__ROCM)
 
@@ -234,10 +234,10 @@ double Memory::record_gpu
 {
 	if(!Memory::init_flag_gpu)
 	{
-		name_gpu = new std::string[n_memory];
-		class_name_gpu = new std::string[n_memory];
-		consume_gpu = new double[n_memory];
-		for(int i=0;i<n_memory;i++)
+		name_gpu = new std::string[n_memory + 2];
+		class_name_gpu = new std::string[n_memory + 2];
+		consume_gpu = new double[n_memory + 2];
+		for(int i=0;i<n_memory + 2;i++)
 		{
 			consume_gpu[i] = 0.0;
 		}
@@ -365,18 +365,15 @@ void Memory::finish(std::ofstream &ofs)
 		delete[] name_gpu;
 		delete[] class_name_gpu;
 		delete[] consume_gpu;
+		init_flag_gpu = false;
 	}
 #endif
 	return;
 }
 
 void Memory::print_all(std::ofstream &ofs)
 {
-	if(!init_flag
-#if defined(__CUDA) || defined(__ROCM)
-	&& !init_flag_gpu
-#endif
-	) 
+	if(!init_flag) 
 	{
 		return;
 	}
@@ -437,6 +434,11 @@ void Memory::print_all(std::ofstream &ofs)
 	}
 
 #if defined(__CUDA) || defined(__ROCM)
+	if(!init_flag_gpu) 
+	{
+		return;
+	}
+
 	ofs <<"\n NAME-------------------------|GPU MEMORY(MB)----" << std::endl;
 	ofs <<std::setw(30)<< "total" << std::setw(15) <<std::setprecision(4)<< Memory::total_gpu << std::endl;