File size: 464 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | """Shared helpers for backbones."""
from __future__ import annotations
import torch
def _pack_scalar(s: torch.Tensor, B: int) -> torch.Tensor:
"""Normalize a scalar / per-batch scalar to shape (B, 1)."""
if s.ndim == 0 or (s.ndim == 1 and s.numel() == 1):
return s.view(1, 1).expand(B, 1)
return s.reshape(B, 1)
def count_parameters(model: torch.nn.Module) -> int:
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|