1 files changed, 5 insertions, 2 deletions
diff --git a/compiler/moco-tf/src/Canonicalization/ReluCanonicalizer.h b/compiler/moco-tf/src/Canonicalization/ReluCanonicalizer.h
index 97adba308..e27abe158 100644
--- a/compiler/moco-tf/src/Canonicalization/ReluCanonicalizer.h
+++ b/compiler/moco-tf/src/Canonicalization/ReluCanonicalizer.h
@@ -18,6 +18,9 @@
 #define __MOCO_TF_RELU_CANONICALIZER_H__
 
 #include "Transform.h"
+#include "SimpleNodeTransform.h"
+
+#include <moco/IR/TFNodes.h>
 
 #include <loco.h>
 
@@ -29,13 +32,13 @@ namespace tf
 /**
  * @brief  Convert TFRelu to Canonical ReLU
  */
-class ReluCanonicalizer : public Transform
+class ReluCanonicalizer : public SimpleNodeTransform<moco::TFRelu>
 {
 public:
   const char *name(void) const final { return "ReluCanonicalizer"; }
 
 public:
-  bool run(loco::Graph *graph) override;
+  bool transform(moco::TFRelu *) const final;
 };
 
 } // namespace tf