diff --git a/comms/torchcomms/rcclx/tests/unit/cpp/mocks/CachingAllocatorHookMock.hpp b/comms/torchcomms/rcclx/tests/unit/cpp/mocks/CachingAllocatorHookMock.hpp
index dc8a621cb..6219f4af0 100644
--- a/comms/torchcomms/rcclx/tests/unit/cpp/mocks/CachingAllocatorHookMock.hpp
+++ b/comms/torchcomms/rcclx/tests/unit/cpp/mocks/CachingAllocatorHookMock.hpp
@@ -11,11 +11,7 @@ namespace torch::comms::test {
 
 class CachingAllocatorHookMock : public CachingAllocatorHookImpl {
  public:
-  MOCK_METHOD(
-      void,
-      regDeregMem,
-      (const c10::hip::HIPCachingAllocator::TraceEntry& entry),
-      (override));
+  MOCK_METHOD(void, regDeregMem, (const TraceEntry& entry), (override));
   MOCK_METHOD(void, registerComm, (TorchCommRCCLX * comm), (override));
   MOCK_METHOD(void, deregisterComm, (TorchCommRCCLX * comm), (override));
   MOCK_METHOD(void, clear, (), (override));