fix import and some error

agents/policy_gradient_agent.py

@@ -6,7 +6,8 @@ from torch.optim import SGD
 
 from tugo.agents.base import Agent
 from tugo.agents.helpers import is_point_an_eye
-from game_logic import goboard
+# from tugo.game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 
 
 def policy_gradient_loss(y_true, y_pred):
@@ -26,7 +27,8 @@ class PolicyAgent(Agent):
 
     def predict(self, game_state):
         encoded_state = self._encoder.encode(game_state)
-        input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
+        # input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
+        input_tensor = encoded_state.float().to('cuda')
         with torch.no_grad():
             output_tensor = self._model(input_tensor)
         return output_tensor.cpu().numpy()[0]

agents/predict.py

@@ -1,6 +1,7 @@
 from tugo.agents.base import Agent
 from tugo.agents.helpers import is_point_an_eye
-from game_logic import goboard
+# from game_logic import goboard
+from tugo.game_logic import goboard_fast as goboard
 from tugo.models import AlphaGoModel
 
 import numpy as np
@@ -44,6 +45,7 @@ class DeepLearningAgent(Agent):
     @classmethod
     def load(cls, file_path, encoder):
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
+        # model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
+        model = AlphaGoModel(encoder.shape(), is_policy_net=True, num_classes=361).to(device)
         model.load_state_dict(torch.load(file_path))
         return cls(model, encoder)