Pokemon Model

model.py

import torch.nn as nn

class PokeAgent(nn.Module):
    def __init__(self, vocab_size, num_choices):
        super(PokeAgent, self).__init__()
        
        self.embedding = nn.Embedding(vocab_size, 32)
        self.linear = nn.Linear(32, 32)
        self.activation = nn.ReLU()
        self.out = nn.Linear(32, num_choices)
    
    def forward(self, turn_input):
        embedding = self.embedding(turn_input)
        embedding = embedding.mean(dim=0)
        
        x = self.linear(embedding)
        x = self.activation(x)
        x = self.out(x)
        
        return x