def forward(self, x):
if self.expk is None or self.expk.size(-1) != x.size(-1):
self.expk = torch.empty(x.size(-1), dtype=x.dtype, device=x.device)
if x.is_cuda:
dct_cuda.precompute_idct_cos(x.size(-1), self.expk)
else:
dct_cpp.precompute_idct_cos(x.size(-1), self.expk)
if self.out is None or self.out.size() != x.size():
self.buf = torch.empty_like(x)
self.out = torch.empty_like(x)
return IDCTFunction.apply(x, self.expk, self.buf, self.out)
def forward(self, x):
if self.expk0 is None or self.expk0.size(-1) != x.size(-2):
self.expk0 = torch.empty(x.size(-2),
dtype=x.dtype,
device=x.device)
if x.is_cuda:
dct_cuda.precompute_idct_cos(x.size(-2), self.expk0)
else:
dct_cpp.precompute_idct_cos(x.size(-2), self.expk0)
if self.expk1 is None or self.expk1.size(-1) != x.size(-1):
self.expk1 = torch.empty(x.size(-2),
dtype=x.dtype,
device=x.device)
if x.is_cuda:
dct_cuda.precompute_idct_cos(x.size(-1), self.expk1)
else:
dct_cpp.precompute_idct_cos(x.size(-1), self.expk1)
if self.out is None or self.out.size() != x.size():
self.buf0 = torch.empty_like(x)
self.buf1 = torch.empty_like(x)
self.out = torch.empty_like(x)
return IDSCT2Function.apply(x, self.expk0, self.expk1, self.buf0,
self.buf1, self.out)
