@MrPositron  Yes, for sure. I define all_params as follows:
# Model initialization
model = LinkPredict()
embed_layer = EmbedLayer()

# Parameter initialization
learning_rate: float = 0.01
weight_decay: float = 0.0

all_params = itertools.chain(model.parameters(), embed_layer.parameters())
optimizer = th.optim.Adam(all_params, lr=learning_rate, weight_decay=weight_decay)

I simplified the model initialization, but the model is the LinkPredict as mentioned above and the EmbedLayer is defined as follows:
class EmbedLayer(nn.Module):
    r"""Embedding layer.."""

    def __init__(self,
                 ...
):
        super(EmbedLayer, self).__init__()