ertk.pytorch.models.get_pt_model_fn

ertk.pytorch.models.get_pt_model_fn(name: str, **kwargs) → Callable[[...], Module]

Get a PyTorch model creation function.

Parameters:

name: str

The model name. This should be the name of the module that implements the model.

**kwargs:

Keyword arguments to provide to the model creation function.

Returns:

model_fn: callable

A method that takes arguments and returns a Module instance.