Class: TorchRec::Models::DLRM::DenseArch

Inherits:

Torch::NN::Module

Object
Torch::NN::Module
TorchRec::Models::DLRM::DenseArch

Defined in:: lib/torchrec/models/dlrm/dense_arch.rb

Instance Method Summary collapse

#forward(features) ⇒ Object
#initialize(in_features, layer_sizes, device: nil) ⇒ DenseArch constructor

A new instance of DenseArch.

Constructor Details

#initialize(in_features, layer_sizes, device: nil) ⇒ `DenseArch`

Returns a new instance of DenseArch.

# File 'lib/torchrec/models/dlrm/dense_arch.rb', line 5

def initialize(in_features, layer_sizes, device: nil)
  super()
  @model = Modules::MLP::MLP.new(
    in_features, layer_sizes, bias: true, activation: :relu, device: device
  )
end

Instance Method Details

#forward(features) ⇒ `Object`



12
13
14

# File 'lib/torchrec/models/dlrm/dense_arch.rb', line 12

def forward(features)
  @model.call(features)
end