forward.py

import torch.nn as nn

def myModule(nn.Module):
  def __init__(self):
    # Init stuff here
    self.X = nn.Sequential(
               nn.Linear(num_input_genes, num_tfs),
               nn.ReLU(),
               nn.BatchNorm1d(num_tfs)
             )
    self.C = nn.Sequential(
               nn.Conv1d(num_tfs, num_conv_out_channels, conv_kernel_size),
               nn.ReLU(),
               nn.BatchNorm1d(num_conv_out_channels),
               nn.MaxPool1d(max_pool_kernel_size)
             )

  def forward(self, input, M):
    x_out = self.X(input)
    x_out = M * x_out # With required reshaping, ...
    x_out = self.C(x_out)
    return x_out