diff --git a/muzero/config.py b/muzero/config.py
index 889f014..477d752 100644
--- a/muzero/config.py
+++ b/muzero/config.py
@@ -1,7 +1,7 @@
 import collections
 from typing import Optional, Dict
 
-import tensorflow_core as tf
+import tensorflow as tf
 
 from game.cartpole import CartPole
 from game.game import AbstractGame
diff --git a/muzero/networks/cartpole_network.py b/muzero/networks/cartpole_network.py
index 8763ad5..34e323c 100644
--- a/muzero/networks/cartpole_network.py
+++ b/muzero/networks/cartpole_network.py
@@ -1,9 +1,9 @@
 import math
 
 import numpy as np
-from tensorflow_core.python.keras import regularizers
-from tensorflow_core.python.keras.layers.core import Dense
-from tensorflow_core.python.keras.models import Sequential
+from tensorflow.keras import regularizers
+from tensorflow.keras.layers import Dense
+from tensorflow.keras.models import Sequential
 
 from game.game import Action
 from networks.network import BaseNetwork
diff --git a/muzero/networks/network.py b/muzero/networks/network.py
index 8912f2c..1b9d60c 100644
--- a/muzero/networks/network.py
+++ b/muzero/networks/network.py
@@ -3,7 +3,7 @@
 from typing import Dict, List, Callable
 
 import numpy as np
-from tensorflow_core.python.keras.models import Model
+from tensorflow.keras.models import Model
 
 from game.game import Action
 
diff --git a/muzero/networks/shared_storage.py b/muzero/networks/shared_storage.py
index c5998da..4d6d593 100644
--- a/muzero/networks/shared_storage.py
+++ b/muzero/networks/shared_storage.py
@@ -1,4 +1,4 @@
-import tensorflow_core as tf
+import tensorflow as tf
 
 from networks.network import BaseNetwork, UniformNetwork, AbstractNetwork
 
diff --git a/muzero/training/training.py b/muzero/training/training.py
index fa6e00f..c80d16f 100644
--- a/muzero/training/training.py
+++ b/muzero/training/training.py
@@ -1,8 +1,8 @@
 """Training module: this is where MuZero neurons are trained."""
 
 import numpy as np
-import tensorflow_core as tf
-from tensorflow_core.python.keras.losses import MSE
+import tensorflow as tf
+from tensorflow.keras.losses import MSE
 
 from config import MuZeroConfig
 from networks.network import BaseNetwork