Is the sparse embedding layer broken (in the master version)? A simple forward and backward pass through the sparse embedding layer seems to have GPU memory issues. It does not go OOM but it keeps on increasing until it reaches the maximum GPU memory and then resets it automatically back to 0. Any idea why that’s happening? (edited)
I am using SparseAdam to update the params for the sparse emb layer.