Class: TorchRec::Modules::MLP::Perceptron

Inherits:
Torch::NN::Module
  • Object
show all
Defined in:
lib/torchrec/modules/mlp/perceptron.rb

Instance Method Summary collapse

Constructor Details

#initialize(in_size, out_size, bias: true, activation: Torch.method(:relu), device: nil) ⇒ Perceptron

Returns a new instance of Perceptron.



5
6
7
8
9
10
11
12
13
14
# File 'lib/torchrec/modules/mlp/perceptron.rb', line 5

def initialize(in_size, out_size, bias: true, activation: Torch.method(:relu), device: nil)
  super()
  @out_size = out_size
  @in_size = in_size
  @linear = Torch::NN::Linear.new(
    # TODO add device
    @in_size, @out_size, bias: bias #, device: device
  )
  @activation_fn = activation
end

Instance Method Details

#forward(input) ⇒ Object



16
17
18
# File 'lib/torchrec/modules/mlp/perceptron.rb', line 16

def forward(input)
  @activation_fn.call(@linear.call(input))
end