Function round_ste

ioblocks.py:127–130

Round with straight through gradients.

(z: Tensor)

# tensor helpers

def round_ste(z: Tensor) -> Tensor:
    """Round with straight through gradients."""
    zhat = z.round()
    return z + (zhat - z).detach()

# main class
# lucidrains fsq