def forward(self, x: torch.Tensor, dt: float = 1.0) -> torch.Tensor
Propagate ``x`` ``(N, D)`` one step and renormalize to unit norm.
Source: unitrack/benchmarks/hota/learned_modules.py:36
