diff --git a/rocket_twin/systems/rocket/stage.py b/rocket_twin/systems/rocket/stage.py
index f0ceceb..3f53f13 100644
--- a/rocket_twin/systems/rocket/stage.py
+++ b/rocket_twin/systems/rocket/stage.py
@@ -2,10 +2,25 @@
 
 from rocket_twin.systems import ControllerCoSApp, Engine, Tank
 
+
 class Stage(System):
+    """Model of a rocket stage.
+
+    Inputs
+    ------
+
+    Outputs
+    ------
+    force [N]: float,
+        thrust force
+    weight [kg]: float,
+        weight
+    cg [m]: float,
+        center of gravity
+    """
 
     def setup(self):
-        
+
         self.add_child(ControllerCoSApp("controller"))
         self.add_child(Tank("tank"), pulling=["w_in", "weight_max", "weight_p"])
         self.add_child(Engine("engine"))
@@ -13,10 +28,12 @@ def setup(self):
         self.connect(self.controller.outwards, self.tank.inwards, {"w": "w_command"})
         self.connect(self.tank.outwards, self.engine.inwards, {"w_out": "w_out"})
 
-        self.add_outward("weight", 1., desc="Weight", unit='kg')
-        self.add_outward("cg", 1., desc="Center of gravity", unit='m')
+        self.add_outward("weight", 1.0, desc="Weight", unit="kg")
+        self.add_outward("cg", 1.0, desc="Center of gravity", unit="m")
 
     def compute(self):
 
         self.weight = self.tank.weight + self.engine.weight
-        self.cg = (self.tank.cg*self.tank.weight + self.engine.cg*self.engine.weight)/(self.weight)
+        self.cg = (self.tank.cg * self.tank.weight + self.engine.cg * self.engine.weight) / (
+            self.weight
+        )