(
schedule,
n_timestep,
linear_start=1e-4,
linear_end=2e-2,
)
| 18 | |
| 19 | |
| 20 | def make_beta_schedule( |
| 21 | schedule, |
| 22 | n_timestep, |
| 23 | linear_start=1e-4, |
| 24 | linear_end=2e-2, |
| 25 | ): |
| 26 | if schedule == "linear": |
| 27 | betas = torch.linspace(linear_start**0.5, linear_end**0.5, n_timestep, dtype=torch.float64) ** 2 |
| 28 | return betas.numpy() |
| 29 | |
| 30 | |
| 31 | def extract_into_tensor(a, t, x_shape): |