diff --git a/ch05/01_main-chapter-code/tests.py b/ch05/01_main-chapter-code/tests.py
index dea2300..e410169 100644
--- a/ch05/01_main-chapter-code/tests.py
+++ b/ch05/01_main-chapter-code/tests.py
@@ -6,7 +6,7 @@
 # File for internal use (unit tests)
 
 import pytest
-from train import main
+from gpt_train import main
 
 
 @pytest.fixture