diff --git a/y2018/control_loops/python/polydrivetrain.py b/y2018/control_loops/python/polydrivetrain.py
index e08d3b6..1652b99 100644
--- a/y2018/control_loops/python/polydrivetrain.py
+++ b/y2018/control_loops/python/polydrivetrain.py
@@ -19,10 +19,10 @@
 def main(argv):
   if FLAGS.plot:
     polydrivetrain.PlotPolyDrivetrainMotions(drivetrain.kDrivetrain)
-  elif len(argv) != 5:
+  elif len(argv) != 7:
     glog.fatal('Expected .h file name and .cc file name')
   else:
-    polydrivetrain.WritePolyDrivetrain(argv[1:3], argv[3:5], 'y2018',
+    polydrivetrain.WritePolyDrivetrain(argv[1:3], argv[3:5], argv[5:7], 'y2018',
                                        drivetrain.kDrivetrain)
 
 if __name__ == '__main__':