Method __init__

diff2flow/models/unet/openaimodel.py:417–419

(self, dim)

class Timestep(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        return timestep_embedding(t, self.dim)