code: import torch.nn as nn
  from chemprop.nn_utils import get_activation_function
  from chemprop.args import TrainArgs
  
  
  def create_ffn(output_size: int, input_size: int, args: TrainArgs):
      """
      Creates the feed-forward layers for the model.
  
      :param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.
      """
  
      first_linear_dim = args.hidden_size * args.number_of_molecules
      # need to also add other 2 network outputs
      if args.use_input_features:
          first_linear_dim += args.features_size
  
      if args.atom_descriptors == "descriptor":
          first_linear_dim += args.atom_descriptors_size
  
      first_linear_dim = input_size
      dropout = nn.Dropout(args.dropout)
      activation = get_activation_function(args.activation)
  
      # Create FFN layers
      if args.ffn_num_layers == 1:
          ffn = [dropout, nn.Linear(first_linear_dim, output_size)]
      else:
          ffn = [dropout, nn.Linear(first_linear_dim, args.ffn_hidden_size)]
          for _ in range(args.ffn_num_layers - 2):
              ffn.extend(
                  [activation, dropout, nn.Linear(args.ffn_hidden_size, args.ffn_hidden_size),]
              )
          ffn.extend(
              [activation, dropout, nn.Linear(args.ffn_hidden_size, output_size),]
          )
  
      # return FFN model
      return nn.Sequential(*ffn)
  
apis: torch.nn.Linear torch.nn.Dropout torch.nn.Sequential