Method __init__

eval_code/recons/models/fastmodel.py:11–21

(
            self,
            pretrained_model_name_or_path: Optional[str] = None,
        )

class Pi3(nn.Module):
    def __init__(
            self,
            pretrained_model_name_or_path: Optional[str] = None,
        ):
        super().__init__()

        if pretrained_model_name_or_path is not None:
            from models.pi3.models.pi3 import Pi3 as Pi3Model
            self.model = Pi3Model.from_pretrained(pretrained_model_name_or_path)
        else:
            raise NotImplementedError

    def forward(self, images: torch.Tensor):
        return self.model(images)