Fixed imports to run with tensorflow 2.3.1

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

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

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
 

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
 

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