tensorless 0.3.0__tar.gz → 0.4.0__tar.gz

{tensorless-0.3.0/tensorless.egg-info → tensorless-0.4.0}/PKG-INFO

@@ -1,20 +1,20 @@
 Metadata-Version: 2.4
 Name: tensorless
-Version: 0.3.0
+Version: 0.4.0
 Summary: ML with maximum automation and minimum setup.
 Author: Tensorless Contributors
 License: MIT
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: torch>=2.0
+Requires-Dist: numpy>=1.24
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Dynamic: license-file
 
 # Tensorless
 
-Tensorless trains small PyTorch models with sensible defaults. It supports
+Tensorless trains small native NumPy models with sensible defaults. It supports
 text generation, text classification, tabular classification, and regression.
 
 ## Install
@@ -37,7 +37,7 @@ tokenizer; use `tokenizer="char"` for a character-level model. Tensorless
 derives model size, batch size, epochs, validation, device, and BPE vocabulary
 size from the data, while every setting can be overridden.
 
-Long text is tokenized lazily and fed through PyTorch in fixed-size batches.
+Long text is tokenized lazily and fed through the native engine in fixed-size batches.
 The automatic batch size uses a token budget; reduce `batch_size` if your
 available memory is limited.
 

{tensorless-0.3.0 → tensorless-0.4.0}/README.md

@@ -1,6 +1,6 @@
 # Tensorless
 
-Tensorless trains small PyTorch models with sensible defaults. It supports
+Tensorless trains small native NumPy models with sensible defaults. It supports
 text generation, text classification, tabular classification, and regression.
 
 ## Install
@@ -23,7 +23,7 @@ tokenizer; use `tokenizer="char"` for a character-level model. Tensorless
 derives model size, batch size, epochs, validation, device, and BPE vocabulary
 size from the data, while every setting can be overridden.
 
-Long text is tokenized lazily and fed through PyTorch in fixed-size batches.
+Long text is tokenized lazily and fed through the native engine in fixed-size batches.
 The automatic batch size uses a token budget; reduce `batch_size` if your
 available memory is limited.
 

{tensorless-0.3.0 → tensorless-0.4.0}/docs/api_reference.md

@@ -59,7 +59,7 @@ Returned by both `train()` and `load()`.
 | `.info()` | all | Dict summary: task, model type, versions, config, metrics, param count |
 
 Attributes: `.task`, `.model_type`, `.config`, `.meta`, `.metrics`,
-`.dataset_fingerprint`, `.model` (the underlying `torch.nn.Module`),
+`.dataset_fingerprint`, `.model` (the underlying native Tensorless model),
 `.tokenizer` (`CharTokenizer` or `None`), `.preprocessor`
 (`TabularPreprocessor` or `None`).
 

{tensorless-0.3.0 → tensorless-0.4.0}/docs/architecture.md

@@ -37,7 +37,7 @@ tensorless/
 ├── serialization/
 │   └── tl_format.py                     save_tl/load_tl: the .tl file format
 ├── devices/
-│   └── device.py                         hardware auto-detection + torch.device resolution
+│   └── device.py                         native CPU device resolution
 └── cli/
     └── main.py                            argparse-based CLI
 ```

tensorless-0.4.0/docs/installation.md

@@ -0,0 +1,29 @@
+# Installation
+
+## Requirements
+
+- Python 3.9 or later
+- NumPy 1.24 or later (installed automatically as a dependency)
+- Tensorless currently runs its native vectorized engine on CPU; accelerator
+  backends are planned without changing the public API.
+
+## Verify your installation
+
+```bash
+python -c "import tensorless as tl; print(tl.__version__)"
+tensorless --help
+```
+
+You should see a version string printed and the CLI's help text.
+
+## Runtime support
+
+The native engine uses NumPy vectorized CPU kernels and automatically resolves
+the device to CPU. You can still specify `device="cpu"` explicitly in
+`tl.train(...)`; the device option remains forward-compatible with future
+accelerator backends.
+
+## Troubleshooting installation
+
+See [troubleshooting.md](troubleshooting.md#installation-issues) for
+common installation problems.

{tensorless-0.3.0 → tensorless-0.4.0}/docs/tl_format.md

@@ -6,9 +6,9 @@ original training data or code beyond having Tensorless installed.
 
 ## Structure
 
-Under the hood, `.tl` is a `torch.save`/`torch.load`-compatible pickle
-archive (see `tensorless/serialization/tl_format.py`) containing a
-dictionary:
+Under the hood, `.tl` is a portable Python pickle file written by the native
+Tensorless serializer (see `tensorless/serialization/tl_format.py`) containing
+a dictionary:
 
 ```python
 {
@@ -18,7 +18,7 @@ dictionary:
     "model_type": "transformer",       # or mlp
     "config": { ... },                  # full resolved TrainConfig used
     "meta": { ... },                    # vocab_size / n_classes / column info -- whatever the model needs to rebuild
-    "model_state_dict": { ... },        # PyTorch model weights
+   "model_state_dict": { ... },        # native NumPy model weights
     "tokenizer_state": { ... } | None,  # CharTokenizer vocab, for text tasks
     "preprocessor_state": { ... } | None,  # TabularPreprocessor state, for tabular tasks
     "dataset_fingerprint": "...",       # fingerprint of the training dataset

{tensorless-0.3.0 → tensorless-0.4.0}/docs/troubleshooting.md

@@ -2,11 +2,9 @@
 
 ## Installation issues
 
-**`ModuleNotFoundError: No module named 'torch'`**
-Torch is a dependency and should install automatically with `pip
-install -e .`. If it didn't, install it manually following
-[pytorch.org/get-started](https://pytorch.org/get-started/locally/) for
-your platform, then reinstall Tensorless.
+**`ModuleNotFoundError: No module named 'numpy'`**
+NumPy is installed automatically with `pip install -e .`. If it is missing,
+install it with `python -m pip install numpy`, then reinstall Tensorless.
 
 **`tensorless: command not found` after installing**
 Make sure the Python environment's `bin`/`Scripts` directory is on your
@@ -69,8 +67,8 @@ fingerprint across runs.
 
 **`CheckpointError: Checkpoint at '...' is corrupt or incompatible.`**
 The checkpoint file was likely truncated by an interruption during the
-(non-atomic part of the) write, or created by an incompatible PyTorch
-version. Delete the `.ckpt` directory and retrain — you'll lose progress
+  (non-atomic part of the) write, or created by an incompatible Tensorless
+  version. Delete the `.ckpt` directory and retrain — you'll lose progress
 on the interrupted run, but the checkpoint being unreadable means it
 can't be safely resumed regardless.
 

{tensorless-0.3.0 → tensorless-0.4.0}/pyproject.toml

@@ -4,14 +4,14 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorless"
-version = "0.3.0"
+version = "0.4.0"
 description = "ML with maximum automation and minimum setup."
 readme = "README.md"
 requires-python = ">=3.9"
 license = { text = "MIT" }
 authors = [{ name = "Tensorless Contributors" }]
 dependencies = [
-    "torch>=2.0",
+    "numpy>=1.24",
 ]
 
 [project.optional-dependencies]

{tensorless-0.3.0 → tensorless-0.4.0}/tensorless/auto/config.py

@@ -115,7 +115,7 @@ def resolve_config(ds: Dataset, user: TrainConfig) -> ResolvedConfig:
         tokenizer=tokenizer,
         bpe_vocab_size=user.bpe_vocab_size if user.bpe_vocab_size is not None else _auto_vocab_size(ds),
         optimizer=user.optimizer or "adamw",
-        learning_rate=user.learning_rate or (3e-4 if model_type == "transformer" else 1e-3),
+        learning_rate=user.learning_rate or (3e-3 if task == "text-classification" else (3e-4 if model_type == "transformer" else 1e-3)),
         weight_decay=user.weight_decay if user.weight_decay is not None else 0.01,
         batch_size=user.batch_size if user.batch_size is not None else _auto_batch_size(n, max_seq_len),
         epochs=user.epochs or _auto_epochs(n),

{tensorless-0.3.0 → tensorless-0.4.0}/tensorless/checkpoint/manager.py

@@ -18,7 +18,7 @@ import shutil
 import tempfile
 from typing import Any, Dict, Optional
 
-import torch
+import pickle
 
 from ..errors import CheckpointError
 
@@ -44,7 +44,8 @@ class CheckpointManager:
         fd, tmp_path = tempfile.mkstemp(dir=self.checkpoint_dir, suffix=".tmp")
         os.close(fd)
         try:
-            torch.save(state, tmp_path)
+            with open(tmp_path, "wb") as handle:
+                pickle.dump(state, handle, protocol=pickle.HIGHEST_PROTOCOL)
             shutil.move(tmp_path, self.path)
         except Exception as e:
             if os.path.exists(tmp_path):
@@ -55,7 +56,8 @@ class CheckpointManager:
         if not self.exists():
             raise CheckpointError(f"No checkpoint found at '{self.path}'.")
         try:
-            return torch.load(self.path, map_location=map_location, weights_only=False)
+            with open(self.path, "rb") as handle:
+                return pickle.load(handle)
         except Exception as e:
             raise CheckpointError(
                 f"Checkpoint at '{self.path}' is corrupt or incompatible: {e}"

{tensorless-0.3.0 → tensorless-0.4.0}/tensorless/data/tabular.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Tuple
 
-import torch
+import numpy as np
 
 _MISSING = "<missing>"
 _UNK = "<unk>"
@@ -132,13 +132,13 @@ class TabularPreprocessor:
 
     def transform(
         self, records: List[Dict[str, Any]], with_target: bool = True
-    ) -> Dict[str, torch.Tensor]:
+    ) -> Dict[str, np.ndarray]:
         n = len(records)
         num_cols = self.numeric_columns
         cat_cols = self.categorical_columns
 
-        numeric = torch.zeros(n, max(len(num_cols), 1), dtype=torch.float32)
-        categorical = torch.zeros(n, max(len(cat_cols), 1), dtype=torch.long)
+        numeric = np.zeros((n, max(len(num_cols), 1)), dtype=np.float32)
+        categorical = np.zeros((n, max(len(cat_cols), 1)), dtype=np.int64)
 
         for i, r in enumerate(records):
             for j, col in enumerate(num_cols):
@@ -158,14 +158,14 @@ class TabularPreprocessor:
 
         if with_target and self.target_column is not None:
             if self.task == "regression":
-                target = torch.zeros(n, dtype=torch.float32)
+                target = np.zeros(n, dtype=np.float32)
                 for i, r in enumerate(records):
                     v = _try_float(r.get(self.target_column))
                     v = self.target_mean if v is None else v
                     target[i] = (v - self.target_mean) / self.target_std
                 out["target"] = target
             else:
-                target = torch.zeros(n, dtype=torch.long)
+                target = np.zeros(n, dtype=np.int64)
                 for i, r in enumerate(records):
                     raw = str(r.get(self.target_column))
                     idx = self.classes.index(raw) if raw in self.classes else 0
@@ -174,10 +174,10 @@ class TabularPreprocessor:
 
         return out
 
-    def inverse_target(self, values: torch.Tensor) -> List[Any]:
+    def inverse_target(self, values: np.ndarray) -> List[Any]:
         if self.task == "regression":
-            return [(v.item() * self.target_std + self.target_mean) for v in values]
-        return [self.classes[int(v.item())] for v in values]
+            return [float(v * self.target_std + self.target_mean) for v in values]
+        return [self.classes[int(v)] for v in values]
 
     def categorical_vocab_sizes(self) -> List[int]:
         return [len(self.column_stats[c].vocab) for c in self.categorical_columns]

tensorless-0.4.0/tensorless/devices/__init__.py

@@ -0,0 +1,3 @@
+from .device import auto_select_device, get_device
+
+__all__ = ["auto_select_device", "get_device"]

{tensorless-0.3.0 → tensorless-0.4.0}/tensorless/devices/device.py

@@ -11,37 +11,21 @@ from __future__ import annotations
 
 from typing import Optional, Tuple
 
-import torch
-
 
 def _tpu_available() -> bool:
-    try:
-        import torch_xla.core.xla_model as xm  # noqa: F401
-
-        return True
-    except Exception:
-        return False
+    return False
 
 
 def _cuda_available() -> bool:
-    try:
-        return torch.cuda.is_available() and torch.cuda.device_count() > 0
-    except Exception:
-        return False
+    return False
 
 
 def _mps_available() -> bool:
-    try:
-        return torch.backends.mps.is_available()
-    except Exception:
-        return False
+    return False
 
 
 def _cuda_supports_bf16() -> bool:
-    try:
-        return torch.cuda.is_bf16_supported()
-    except Exception:
-        return False
+    return False
 
 
 def auto_select_device(user_device: Optional[str], user_precision: Optional[str]) -> Tuple[str, str]:
@@ -85,23 +69,8 @@ def auto_select_device(user_device: Optional[str], user_precision: Optional[str]
     return device, precision
 
 
-def get_torch_device(device: str) -> torch.device:
-    """Convert our string device name into a torch.device, with a
-    runtime fallback to CPU if the requested backend is unavailable.
-    """
-    try:
-        if device == "tpu":
-            import torch_xla.core.xla_model as xm
-
-            return xm.xla_device()
-        if device == "cuda":
-            if not _cuda_available():
-                return torch.device("cpu")
-            return torch.device("cuda")
-        if device == "mps":
-            if not _mps_available():
-                return torch.device("cpu")
-            return torch.device("mps")
-        return torch.device("cpu")
-    except Exception:
-        return torch.device("cpu")
+def get_device(device: str) -> str:
+    """Return the native engine device name (currently CPU only)."""
+    return "cpu"
+
+

tensorless-0.4.0/tensorless/engine.py

@@ -0,0 +1,187 @@
+"""Small NumPy training engine used by Tensorless models."""
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, Iterator
+import numpy as np
+
+
+class Parameter:
+    def __init__(self, data):
+        self.data = np.asarray(data, dtype=np.float32)
+        self.grad = np.zeros_like(self.data)
+
+    def numel(self) -> int:
+        return int(self.data.size)
+
+
+class Module:
+    def parameters(self) -> Iterator[Parameter]:
+        for value in self.__dict__.values():
+            if isinstance(value, Parameter):
+                yield value
+            elif isinstance(value, Module):
+                yield from value.parameters()
+            elif isinstance(value, (list, tuple)):
+                for item in value:
+                    if isinstance(item, Parameter):
+                        yield item
+                    elif isinstance(item, Module):
+                        yield from item.parameters()
+
+    def named_parameters(self, prefix: str = ""):
+        for name, value in self.__dict__.items():
+            full = f"{prefix}.{name}" if prefix else name
+            if isinstance(value, Parameter):
+                yield full, value
+            elif isinstance(value, Module):
+                yield from value.named_parameters(full)
+            elif isinstance(value, (list, tuple)):
+                for i, item in enumerate(value):
+                    if isinstance(item, Parameter):
+                        yield f"{full}.{i}", item
+                    elif isinstance(item, Module):
+                        yield from item.named_parameters(f"{full}.{i}")
+
+    def state_dict(self) -> Dict[str, np.ndarray]:
+        return {name: parameter.data.copy() for name, parameter in self.named_parameters()}
+
+    def load_state_dict(self, state: Dict[str, np.ndarray]) -> None:
+        current = dict(self.named_parameters())
+        missing = [name for name in current if name not in state]
+        if missing:
+            raise ValueError(f"Missing model parameters: {missing}")
+        for name, parameter in current.items():
+            value = np.asarray(state[name], dtype=np.float32)
+            if value.shape != parameter.data.shape:
+                raise ValueError(f"Shape mismatch for {name}: {value.shape} != {parameter.data.shape}")
+            parameter.data[...] = value
+
+    def train(self, mode: bool = True):
+        self.training = mode
+        return self
+
+    def eval(self):
+        return self.train(False)
+
+
+def softmax_cross_entropy(logits: np.ndarray, targets: np.ndarray, ignore_index=None):
+    flat = logits.reshape(-1, logits.shape[-1]).astype(np.float64)
+    target = targets.reshape(-1).astype(np.int64)
+    valid = np.ones(len(target), dtype=bool) if ignore_index is None else target != ignore_index
+    if not np.any(valid):
+        return 0.0, np.zeros_like(logits)
+    shifted = flat - flat.max(axis=1, keepdims=True)
+    probs = np.exp(shifted)
+    probs /= probs.sum(axis=1, keepdims=True)
+    rows = np.arange(len(target))[valid]
+    loss = -np.log(np.maximum(probs[rows, target[valid]], 1e-12)).mean()
+    grad = np.zeros_like(flat, dtype=np.float32)
+    grad[valid] = probs[valid].astype(np.float32)
+    grad[rows, target[valid]] -= 1.0
+    grad[valid] /= valid.sum()
+    return float(loss), grad.reshape(logits.shape)
+
+
+class Optimizer:
+    def __init__(self, parameters: Iterable[Parameter], lr: float, weight_decay: float = 0.0):
+        self.parameters_list = list(parameters)
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.step_count = 0
+
+    def zero_grad(self):
+        for parameter in self.parameters_list:
+            parameter.grad.fill(0.0)
+
+    def state_dict(self):
+        return {"step": self.step_count, "lr": self.lr, "weight_decay": self.weight_decay}
+
+    def load_state_dict(self, state):
+        self.step_count = int(state.get("step", 0))
+        self.lr = float(state.get("lr", self.lr))
+        self.weight_decay = float(state.get("weight_decay", self.weight_decay))
+
+
+class SGD(Optimizer):
+    def __init__(self, parameters, lr, weight_decay=0.0, momentum=0.9):
+        super().__init__(parameters, lr, weight_decay)
+        self.momentum = momentum
+        self.velocity = [np.zeros_like(p.data) for p in self.parameters_list]
+
+    def step(self):
+        self.step_count += 1
+        for p, velocity in zip(self.parameters_list, self.velocity):
+            velocity *= self.momentum
+            velocity += p.grad + self.weight_decay * p.data
+            p.data -= self.lr * velocity
+
+    def state_dict(self):
+        state = super().state_dict()
+        state["velocity"] = [v.copy() for v in self.velocity]
+        return state
+
+    def load_state_dict(self, state):
+        super().load_state_dict(state)
+        for target, source in zip(self.velocity, state.get("velocity", [])):
+            target[...] = source
+
+
+class Adam(Optimizer):
+    def __init__(self, parameters, lr, weight_decay=0.0, decoupled=False):
+        super().__init__(parameters, lr, weight_decay)
+        self.decoupled = decoupled
+        self.m = [np.zeros_like(p.data) for p in self.parameters_list]
+        self.v = [np.zeros_like(p.data) for p in self.parameters_list]
+
+    def step(self):
+        self.step_count += 1
+        for p, m, v in zip(self.parameters_list, self.m, self.v):
+            if self.decoupled:
+                p.data *= 1.0 - self.lr * self.weight_decay
+                gradient = p.grad
+            else:
+                gradient = p.grad + self.weight_decay * p.data
+            m *= 0.9
+            m += 0.1 * gradient
+            v *= 0.999
+            v += 0.001 * gradient * gradient
+            m_hat = m / (1.0 - 0.9 ** self.step_count)
+            v_hat = v / (1.0 - 0.999 ** self.step_count)
+            p.data -= self.lr * m_hat / (np.sqrt(v_hat) + 1e-8)
+
+    def state_dict(self):
+        state = super().state_dict()
+        state.update({"m": [x.copy() for x in self.m], "v": [x.copy() for x in self.v]})
+        return state
+
+    def load_state_dict(self, state):
+        super().load_state_dict(state)
+        for target, source in zip(self.m, state.get("m", [])):
+            target[...] = source
+        for target, source in zip(self.v, state.get("v", [])):
+            target[...] = source
+
+
+class LambdaScheduler:
+    def __init__(self, optimizer, warmup_steps, total_steps):
+        self.optimizer = optimizer
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.step_count = 0
+        self.base_lr = optimizer.lr
+
+    def step(self):
+        self.step_count += 1
+        step = self.step_count - 1
+        if self.warmup_steps and step < self.warmup_steps:
+            factor = (step + 1) / self.warmup_steps
+        else:
+            factor = max(0.1, 1.0 - (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps))
+        self.optimizer.lr = self.base_lr * factor
+
+    def state_dict(self):
+        return {"step_count": self.step_count}
+
+    def load_state_dict(self, state):
+        self.step_count = int(state.get("step_count", 0))

tensorless-0.4.0/tensorless/models/mlp.py

@@ -0,0 +1,71 @@
+"""NumPy MLP for tabular classification and regression."""
+
+from __future__ import annotations
+
+from typing import List
+import numpy as np
+
+from ..engine import Module, Parameter
+
+
+def _embedding_dim(vocab_size: int) -> int:
+    return max(2, min(32, round(1.6 * (vocab_size ** 0.56))))
+
+
+class TabularMLP(Module):
+    def __init__(self, n_numeric, categorical_vocab_sizes: List[int], d_model, layers, dropout, task, n_classes=0):
+        self.training = True
+        self.task = task
+        self.n_numeric = n_numeric
+        self.categorical_vocab_sizes = categorical_vocab_sizes
+        rng = np.random.default_rng()
+        self.embeddings = [Parameter(rng.normal(0, .02, (v, _embedding_dim(v))).astype(np.float32)) for v in categorical_vocab_sizes]
+        in_dim = n_numeric + sum(_embedding_dim(v) for v in categorical_vocab_sizes)
+        self.weights, self.biases = [], []
+        for i in range(layers):
+            fan_in = in_dim if i == 0 else d_model
+            self.weights.append(Parameter(rng.normal(0, .02, (fan_in, d_model)).astype(np.float32)))
+            self.biases.append(Parameter(np.zeros(d_model, dtype=np.float32)))
+        out_dim = n_classes if task == "classification" else 1
+        self.head_weight = Parameter(rng.normal(0, .02, (d_model if layers else in_dim, out_dim)).astype(np.float32))
+        self.head_bias = Parameter(np.zeros(out_dim, dtype=np.float32))
+
+    def forward(self, numeric, categorical, cache=False):
+        parts = [numeric.astype(np.float32)] if self.n_numeric else []
+        for i, embedding in enumerate(self.embeddings):
+            parts.append(embedding.data[categorical[:, i]])
+        x = np.concatenate(parts, axis=1) if parts else numeric.astype(np.float32)
+        activations = [x]
+        for weight, bias in zip(self.weights, self.biases):
+            x = np.maximum(0, x @ weight.data + bias.data)
+            activations.append(x)
+        out = x @ self.head_weight.data + self.head_bias.data
+        if self.task == "regression": out = out[:, 0]
+        return (out, (activations, categorical)) if cache else out
+
+    def loss_and_backward(self, numeric, categorical, target):
+        logits, (activations, categorical) = self.forward(numeric, categorical, cache=True)
+        if self.task == "regression":
+            error = logits - target
+            loss = float(np.mean(error * error))
+            grad = (2.0 / len(target)) * error[:, None]
+        else:
+            shifted = logits - logits.max(1, keepdims=True)
+            probs = np.exp(shifted); probs /= probs.sum(1, keepdims=True)
+            loss = float(-np.log(np.maximum(probs[np.arange(len(target)), target], 1e-12)).mean())
+            grad = probs; grad[np.arange(len(target)), target] -= 1; grad /= len(target)
+        self.head_weight.grad[...] = activations[-1].T @ grad
+        self.head_bias.grad[...] = grad.sum(0)
+        dx = grad @ self.head_weight.data.T
+        for i in range(len(self.weights) - 1, -1, -1):
+            dx = dx * (activations[i + 1] > 0)
+            self.weights[i].grad[...] = activations[i].T @ dx
+            self.biases[i].grad[...] = dx.sum(0)
+            dx = dx @ self.weights[i].data.T
+        offset = self.n_numeric
+        for i, embedding in enumerate(self.embeddings):
+            width = embedding.data.shape[1]
+            embedding.grad.fill(0)
+            np.add.at(embedding.grad, categorical[:, i], dx[:, offset:offset + width])
+            offset += width
+        return loss

{tensorless-0.3.0 → tensorless-0.4.0}/tensorless/models/registry.py

@@ -11,14 +11,12 @@ from __future__ import annotations
 
 from typing import Any, Dict
 
-import torch.nn as nn
-
 from .transformer import TinyTransformer
 from .mlp import TabularMLP
 from ..errors import ModelError
 
 
-def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str, Any]) -> nn.Module:
+def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str, Any]):
     """Build a fresh, randomly-initialized model.
 
     `cfg` is the resolved training config (dict). `meta` carries