Function sample

examples/minrf.py:149–150

(z:Tensor, cond:Tensor)

  @TinyJit
  def sample(z:Tensor, cond:Tensor) -> Tensor:
    return model.sample(z, cond, Tensor.full_like(cond, 10), sample_steps=getenv("SAMPLE_STEPS", 20))[-1]

  for steps in (t:=trange(getenv("STEPS", 5000))):
    if steps%10 == 0: mviz(sample(Tensor.randn(3, 1, 32, 32), Tensor([5,0,4], dtype='int')))