From dda898c0ab33688f53acde61d91a64df4afd0e29 Mon Sep 17 00:00:00 2001
From: Sandeep Kumar <Sandeep.Kumar@amd.com>
Date: Tue, 16 Oct 2018 02:41:58 -0700
Subject: [PATCH] print HIP memory usage

---
 caffe2/python/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/caffe2/python/utils.py b/caffe2/python/utils.py
index 5e87df8058e01..94e370ffbb7c0 100644
--- a/caffe2/python/utils.py
+++ b/caffe2/python/utils.py
@@ -237,7 +237,7 @@ def ConvertProtoToBinary(proto_class, filename, out_filename):
 
 
 def GetGPUMemoryUsageStats():
-    """Get GPU memory usage stats from CUDAContext. This requires flag
+    """Get GPU memory usage stats from CUDAContext/HIPContext. This requires flag
        --caffe2_gpu_memory_tracking to be enabled"""
     from caffe2.python import workspace, core
     workspace.RunOperatorOnce(
@@ -245,7 +245,7 @@ def GetGPUMemoryUsageStats():
             "GetGPUMemoryUsage",
             [],
             ["____mem____"],
-            device_option=core.DeviceOption(caffe2_pb2.CUDA, 0),
+            device_option=core.DeviceOption(caffe2_pb2.CUDA if workspace.has_gpu_support else caffe2_pb2.HIP, 0),
         ),
     )
     b = workspace.FetchBlob("____mem____")