torch_blue.vi.VIResidualConnection

class torch_blue.vi.VIResidualConnection(*args: torch.nn.Module)

class VIResidualConnection(arg: OrderedDict[str, torch.nn.Module])

Bases: VISequential

A version of VISequential that supports residual connections.

This class is identical to VISequential, but adds the input to the output. Importantly, it manages log prob tracking, if required. Note that a single module can also be wrapped to add a residual connection around it.

Raises:

RuntimeError – If the output shape does not match the input shape during the forward pass.

forward(input_)

Forward pass that manages log probs, if required, and adds the input to the output.

Parameters:

input (Varies) – Input for the first module in the stack. Passed on to it unchanged.

Returns:

output – Output of the module stack plus the input to the residual connection.

Return type:

Varies