Providers in Detail

Model Provider

def skyline_model_provider() -> torch.nn.Module:

The model provider must take no arguments and return an instance of your model (a torch.nn.Module) that is on the GPU (i.e. you need to call .cuda() on the module before returning it).

Input Provider

def skyline_input_provider(batch_size: int = 32) -> Tuple:

The input provider must take a single batch_size argument that has a default value (the batch size you want to profile with). It must return an iterable (does not have to be a tuple) that contains the arguments that you would normally pass to your model's forward method. Any Tensors in the returned iterable must be on the GPU (i.e. you need to call .cuda() on them before returning them).

Iteration Provider

def skyline_iteration_provider(model: torch.nn.Module) -> Callable:

The iteration provider must take a single model argument, which will be an instance of your model. This provider must return a callable (e.g., a function) that, when invoked, runs a single training iteration.