amandaodoo's picture
Publicar modelo y Model Card
d3d8e85 verified
Raw History Blame Contribute Delete
519 Bytes
"""Arquitectura del modelo. La usan train.py y las carpetas pytorch-exportar-*.
PyTorch no guarda la clase dentro del archivo de pesos: para cargar model.safetensors
hay que volver a crear el mismo modelo vacío y después cargarle los pesos.
"""
import torch
class RegresionLineal(torch.nn.Module):
def __init__(self, n_entradas=3):
super().__init__()
self.lineal = torch.nn.Linear(n_entradas, 1) # pesos (w) y bias (b)
def forward(self, x):
return self.lineal(x)