MNISTMLP

Classes

class backbone.MNISTMLP.MNISTMLP(input_size, output_size)

Bases: MammothBackbone

Network composed of two hidden layers, each containing 100 ReLU activations. Designed for the MNIST dataset.

forward(x, returnt='out')

Compute a forward pass.

Parameters:

x (Tensor) – input tensor (batch_size, input_size)

Returns:

output tensor (output_size)

Return type:

Tensor

reset_parameters()

Calls the Xavier parameter initialization function.

to(device)