Fixed DLRModel / DLRMBlock to accept bottom_block with dropout on top of MLPs

[gabrielspmoreira]

Goals ⚽

This PR fixes an error raised by DLRModel when we provide in bottom_block=MLPBlock(..., dropout) arg set. For example:

model = mm.DLRMModel(
        schema,
        embedding_dim=10,
        bottom_block=mm.MLPBlock([64, 32], dropout=0.05),
        prediction_tasks=mm.BinaryOutput("click"),
    )

Implementation Details 🚧

This PR fixes that by iterating over bottom_block layers and finding the last MLP layer to check fit its output dim matches embedding dim.

Testing Details 🔍

The test_dlrm_model was updated to test bottom_block=mm.MLPBlock([5, 2], dropout=0.05)

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/models/review/pr-1018