PyTorch

morphocut.torch.PyTorch()[source]

Apply a PyTorch module to the input.

Parameters

  • module (torch.nn.Module) – PyTorch module.

  • input (input, Variable) – Input.

  • device (str or torch.device, optional) – Device.

  • is_batch (bool, optional) – Assume that input is a batch.

  • output_key (optional) – If the module has multiple outputs, output_key selects one of them.

  • pin_memory (bool, optional) – Use pinned memory for faster CPU-GPU transfer. Only applicable for CUDA devices. If None, enabled by default for CUDA devices.

  • pre_transform (callable, optional) – Transformation to apply to the individual input values.

  • autocast (bool, optional) – Enable automatic mixed precision inference to improve performance.

Example

module = ...
with Pipeline() as pipeline:
    input = ...
    output = PyTorch(module, input)