RuntimeError: value cannot be converted to type uint8_t without overflow with masked_fill in GPT-Neo causal mask creation on DirectML. torch.where works. #702

# target_for_masked_fill = causal_mask_clone[:, :, :, :mask_length]
# padding_mask = ((causal_mask_clone[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask_clone.device)) == 0)
# min_dtype_val = torch.finfo(dtype).min

# Problematic line:
# causal_mask_clone[:, :, :, :mask_length] = target_for_masked_fill.masked_fill(
#     padding_mask, min_dtype_val
# )

# target_for_op = causal_mask_clone[:, :, :, :mask_length]
# padding_mask_bool = ((causal_mask_clone[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask_clone.device)) == 0) # Simplified for brevity
# min_dtype_val = torch.finfo(dtype).min

filled_slice = torch.where(
    padding_mask_bool,    # Condition (boolean mask where True means fill)
    min_dtype_val,        # Value if condition is True
    target_for_op         # Value if condition is False (original tensor slice)
)
causal_mask_clone[:, :, :, :mask_length] = filled_slice

import torch
import torch_directml as dml
from transformers import AutoTokenizer, AutoModelForCausalLM
import traceback

# Configuración del dispositivo DirectML
device = dml.device(0) 
print(f"Dispositivo configurado: {device}")

# Carga del modelo y el tokenizer
model_name = "EleutherAI/gpt-neo-1.3B"
print(f"Cargando el modelo {model_name} en {device}...")
try:
    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
    print("Modelo cargado correctamente.")
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    traceback.print_exc()
    exit()

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    print("Tokenizer configurado correctamente.")
except Exception as e:
    print(f"Error al cargar el tokenizer: {e}")
    traceback.print_exc()
    exit()

# Definir prompt
prompt = "La inteligencia artificial es"
print(f"Prompt de entrada: {prompt}")

inputs_on_device = {} 

try:
    print("Tokenizando el input...")
    inputs_cpu = tokenizer(prompt, return_tensors="pt", padding=True)
    
    inputs_on_device['input_ids'] = inputs_cpu['input_ids'].to(device=device, dtype=torch.int32)
    if 'attention_mask' in inputs_cpu:
        inputs_on_device['attention_mask'] = inputs_cpu['attention_mask'].to(device=device, dtype=torch.float32)
    else:
        inputs_on_device['attention_mask'] = torch.ones_like(inputs_on_device['input_ids'], dtype=torch.float32, device=device)
    
    print(f"Tipo de input_ids ANTES de la inferencia (en {device}): {inputs_on_device['input_ids'].dtype}")
    print(f"Tipo de attention_mask ANTES de la inferencia (en {device}): {inputs_on_device['attention_mask'].dtype}")

except Exception as e:
    print(f"Error durante la tokenización o movimiento a dispositivo: {e}")
    traceback.print_exc()
    exit()

try:
    print("Realizando inferencia...")
    with torch.no_grad(): 
        output = model.generate(
            inputs_on_device['input_ids'],
            attention_mask=inputs_on_device['attention_mask'],
            max_length=50,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    print(f"Output tensor (on device {output.device}): {output}")
    decoded_output = tokenizer.decode(output[0].to('cpu'), skip_special_tokens=True)
    print(f"Texto generado: {decoded_output}")
except Exception as e:
    print(f"Error durante la inferencia: {e}")
    traceback.print_exc()