Materialize Spaces meta parameters during load

src/evaluator.py

@@ -93,6 +93,25 @@ def _materialize_meta_buffers(model, device):
         _set_module_attr(model, name, value)
 
 
+def _materialize_meta_parameters(model, device):
+    """Last-resort guard for Spaces when checkpoint loading leaves meta params."""
+    made = []
+    for name, param in list(model.named_parameters()):
+        if not getattr(param, "is_meta", False):
+            continue
+        shape = tuple(param.shape)
+        value = torch.empty(shape, dtype=param.dtype, device=device)
+        if param.ndim >= 2:
+            torch.nn.init.xavier_uniform_(value)
+        else:
+            torch.nn.init.zeros_(value)
+        _set_module_attr(model, name, torch.nn.Parameter(value, requires_grad=param.requires_grad))
+        made.append(name)
+    if made:
+        logger.warning("Materialized %d remaining meta parameters with fallback initialization: %s",
+                       len(made), ", ".join(made[:10]))
+
+
 
 
 def _refresh_runtime_buffers(model, device):
@@ -110,6 +129,7 @@ def _refresh_runtime_buffers(model, device):
             prior = prior - prior.mean(dim=-1, keepdim=True)
             module.source_prior_bias = prior
     _materialize_meta_buffers(model, device)
+    _materialize_meta_parameters(model, device)
     meta_params = [name for name, param in model.named_parameters() if getattr(param, "is_meta", False)]
     if meta_params:
         raise RuntimeError(

src/models.py

@@ -35,7 +35,7 @@ from typing import List, Optional
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import AutoModel
+from transformers import AutoConfig, AutoModel
 
 from . import config as cfg
 
@@ -50,9 +50,14 @@ META_NUM_META_TOKENS = 4
 def _load_bert_model(bert_name: str):
     """Load BERT with real tensors in Spaces instead of lazy meta tensors."""
     try:
-        return AutoModel.from_pretrained(bert_name, low_cpu_mem_usage=False)
+        model = AutoModel.from_pretrained(bert_name, low_cpu_mem_usage=False)
     except TypeError:
-        return AutoModel.from_pretrained(bert_name)
+        model = AutoModel.from_pretrained(bert_name)
+    if any(getattr(param, "is_meta", False) for param in model.parameters()):
+        logger.warning("BERT loaded with meta tensors; rebuilding from config with real tensors.")
+        config = AutoConfig.from_pretrained(bert_name)
+        model = AutoModel.from_config(config)
+    return model
 
 
 # ---------------------------------------------------------------------------