diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 79e2864..cbb4225 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -5,7 +5,6 @@
 #include <cmath>
 #include <memory>
 
-#include "aos/aos_core.h"
 #include "aos/common/logging/logging.h"
 #include "aos/common/queue.h"
 #include "frc971/control_loops/state_feedback_loop.h"
diff --git a/frc971/control_loops/drivetrain/drivetrain_main.cc b/frc971/control_loops/drivetrain/drivetrain_main.cc
index a6efbcd..e3fc306 100644
--- a/frc971/control_loops/drivetrain/drivetrain_main.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_main.cc
@@ -1,6 +1,6 @@
 #include "frc971/control_loops/drivetrain/drivetrain.h"
 
-#include "aos/aos_core.h"
+#include "aos/atom_code/init.h"
 
 int main() {
   ::aos::Init();