def METHOD_NAME(self, x):
    for param in self.model.parameters():
        # 'in' to handle the situation when .to('cuda') results in 'cuda:0' actual device
        assert self.original_device in str(param.device)
    return self.model.METHOD_NAME(x)