Para el siguiente codigo me aparece RuntimeError: CUDA error: invalid configuration argument
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions. : import torch
import torch.nn as nn
import time
from torchtext.datasets import PennTreebank
from torchtext.data.functional import to_map_style_dataset
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator


# 1. Definir el modelo
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_layers)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src = self.embedding(src)
        tgt = self.embedding(tgt)
        x = self.transformer(src, tgt)
        x = self.fc(x)
        return x

# 2. Preparar datos
train_data_raw = to_map_style_dataset(PennTreebank(split='train'))
tokenizer = get_tokenizer("spacy", "en_core_web_sm")

# Agregar tokens especiales (UNK y PAD) al vocabulario
specials = ['<unk>', '<pad>']
vocab = build_vocab_from_iterator((tokenizer(y) for y in train_data_raw), specials=specials)
vocab.set_default_index(vocab['<unk>'])

# Dividir los datos en conjuntos de entrenamiento y validación
train_ratio = 0.8
train_data_size = int(len(train_data_raw) * train_ratio)
train_data = train_data_raw[:train_data_size]
valid_data = train_data_raw[train_data_size:]


# 3. Entrenar el modelo
def generate_pairs(tokens, shift=1):
    source = [tokens[i] for i in range(0, len(tokens) - shift)]
    target = [tokens[i] for i in range(shift, len(tokens))]
    return source, target

def encode_text(text):
    tokens = tokenizer(text)
    return torch.tensor([vocab[token] if token in vocab else vocab['<unk>'] for token in tokens]).to(device)

vocab_size = len(vocab)
d_model = 512
nhead = 4
num_layers = 4
num_epochs = 5
learning_rate = 1e-3
batch_size = 64
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = TransformerModel(vocab_size=vocab_size, d_model=d_model, nhead=nhead, num_layers=num_layers).to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

def train(model, epoch):
    model.train()
    train_loss = 0
    for i, text in enumerate(train_data):
        tokens = encode_text(text)
        if len(tokens) <= 1:
            continue
        source, target = generate_pairs(tokens, shift=1)
        source = torch.tensor(source).unsqueeze(1).to(device)
        target = torch.tensor(target).unsqueeze(1).to(device)
        optimizer.zero_grad()
        output = model(source, target[:-1])
        loss = loss_function(output.view(-1, vocab_size), target[1:].view(-1))
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        if i % 1000 == 0:
            print(f'Epoch: {epoch + 1},\tLoss: {loss.item()}')
    return train_loss / len(train_data)

for epoch in range(num_epochs):
    start_time = time.time()
    train_loss = train(model, epoch)
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Epoch: {epoch + 1},\tAverage Loss: {train_loss},\tTime taken: {elapsed_time} seconds")

# 4. Evaluar el modelo en el conjunto de validación
def evaluate(model, valid_data):
    model.eval()
    valid_loss = 0
    with torch.no_grad():
        for i, text in enumerate(valid_data):
            tokens = encode_text(text)
            if len(tokens) <= 1:
                continue
            source, target = generate_pairs(tokens, shift=1)
            source = torch.tensor(source).unsqueeze(1).to(device)
            target = torch.tensor(target).unsqueeze(1).to(device)
            output = model(source, target[:-1])
            loss = loss_function(output.view(-1, vocab_size), target[1:].view(-1))
            valid_loss += loss.item()
    return valid_loss / len(valid_data)

# 5. Entrenar y evaluar el modelo en cada época
for epoch in range(num_epochs):
    start_time = time.time()
    train_loss = train(model, epoch)
    valid_loss = evaluate(model, valid_data)
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Epoch: {epoch + 1},\tAverage Train Loss: {train_loss},\tAverage Valid Loss: {valid_loss},\tTime taken: {elapsed_time} seconds")

# 6. Guardar el modelo entrenado
torch.save(model.state_dict(), "trained_model.pth")

# 7. Cargar el modelo
def load_model(path, device):
    loaded_model = TransformerModel(vocab_size=vocab_size, d_model=d_model, nhead=nhead, num_layers=num_layers).to(device)
    loaded_model.load_state_dict(torch.load(path))
    loaded_model.eval()
    return loaded_model

loaded_model = load_model("trained_model.pth", device)

# 8. Predecir texto
def decode_output(output):
    probabilities = torch.softmax(output, dim=2)
    predicted_indices = torch.argmax(probabilities, dim=2)
    return " ".join(vocab.lookup_tokens(predicted_indices.cpu().numpy().squeeze().tolist()))

def predict_next_token(input_text, model):
    tokens = encode_text(input_text)
    source, target = generate_pairs(tokens, shift=1)
    source = torch.tensor(source).unsqueeze(1).to(device)
    target = torch.tensor(target).unsqueeze(1).to(device)
    with torch.no_grad():
        output = model(source, target[:-1])
        return decode_output(output)

input_text = "My favorite color is"
predicted_output = predict_next_token(input_text, loaded_model)
print(f"Predicted output: {predicted_output}")