Views
No views yet
input_dim: 6hidden_dim: 256output_dim: 31import torch
2import torch.nn as nn
3from safetensors.torch import load_file
4from huggingface_hub import hf_hub_download
5
6class MLP(nn.Module):
7 def __init__(self, input_dim, hidden_dim, output_dim):
8 super().__init__()
9 self.mlp = nn.Sequential(
10 nn.Linear(input_dim, hidden_dim),
11 nn.LayerNorm(hidden_dim),
12 nn.ReLU(),
13 nn.Linear(hidden_dim, hidden_dim),
14 nn.LayerNorm(hidden_dim),
15 nn.ReLU(),
16 nn.Linear(hidden_dim, output_dim),
17 )
18 def forward(self, x):
19 return self.mlp(x)
20
21path = hf_hub_download("chewwt/dm_qwen4b_emulator", "model.safetensors")
22model = MLP(input_dim=6, hidden_dim=256, output_dim=3)
23model.load_state_dict(load_file(path))
24model.eval()
25
26with torch.no_grad():
27 out = model(torch.randn(1, 6))