diff --git a/core/task/taskclass/class.go b/core/task/taskclass/class.go
index 902741c6..0acf88a9 100644
--- a/core/task/taskclass/class.go
+++ b/core/task/taskclass/class.go
@@ -123,7 +123,6 @@ func (c *Class) UnmarshalYAML(unmarshal func(interface{}) error) (err error) {
 		}
 	}
 	return
-
 }
 
 func (c *Class) MarshalYAML() (interface{}, error) {
@@ -153,14 +152,7 @@ func (c *Class) MarshalYAML() (interface{}, error) {
 		Constraints: c.Constraints,
 		Command:     c.Command,
 	}
-
-	if c.Control.Mode == controlmode.FAIRMQ {
-		aux.Control.Mode = "fairmq"
-	} else if c.Control.Mode == controlmode.BASIC {
-		aux.Control.Mode = "basic"
-	} else {
-		aux.Control.Mode = "direct"
-	}
+	aux.Control.Mode = c.Control.Mode.String()
 
 	return aux, nil
 }