diff --git a/core/api/src/main/java/org/onosproject/net/flow/DefaultTrafficTreatment.java b/core/api/src/main/java/org/onosproject/net/flow/DefaultTrafficTreatment.java
index 6b5af75..37dcd06 100644
--- a/core/api/src/main/java/org/onosproject/net/flow/DefaultTrafficTreatment.java
+++ b/core/api/src/main/java/org/onosproject/net/flow/DefaultTrafficTreatment.java
@@ -17,8 +17,10 @@
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
+
 import org.onlab.packet.IpAddress;
 import org.onlab.packet.MacAddress;
+import org.onlab.packet.MplsLabel;
 import org.onlab.packet.VlanId;
 import org.onosproject.core.GroupId;
 import org.onosproject.net.PortNumber;
@@ -222,12 +224,12 @@
         }
 
         @Override
-        public Builder popMpls(short etherType) {
+        public Builder popMpls(Short etherType) {
             return add(Instructions.popMpls(etherType));
         }
 
         @Override
-        public Builder setMpls(Integer mplsLabel) {
+        public Builder setMpls(MplsLabel mplsLabel) {
             return add(Instructions.modMplsLabel(mplsLabel));
         }