DownsamplingGenerator's batching is inconsistent. Currently, step() returns factor: int and filter: Tensor of shape (1, 1, k, k), and the batch_size is unused. This means it can't generate batched params. Furthermore, the factor is incompatible with generate_dataset.

It should return factor: Tensor of shape (B, 1) and filter: Tensor of shape (B, 1, k, k). physics.Downsampling should also be modified. It should either allow batched physics with different factors or filters, or raise a warning that physics params batch size must be one (or all params are constant in one batch).