diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 640ce56883..6bb4e46d5a 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -42,8 +42,6 @@
 from MaxText import sharding
 from MaxText.common_types import ShardMode
 from MaxText.globals import EPS
-# pylint: disable-next=unused-import
-from maxtext import maxtext_google
 
 from MaxText.gradient_accumulation import gradient_accumulation_loss_and_grad
 from MaxText.vocabulary_tiling import vocab_tiling_linen_loss