class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, device=torch.device('cpu')):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, action_dim)
        self.device = device
        self.to(device)