How do we use this in practice? If the embedding layer is sparse but the rest of the network is dense, do we need to have two different optimizers? Is there any way to do this with param groups.
import time                                                                                                                                                                                                       
import numpy                                                                                                                                                                                                      
import torch                                                                                                                                                                                                      
import torch.nn as nn                                                                                                                                                                                             
import torch.nn.functional as F                                                                                                                                                                                   
import torch.optim as optim                                                                                                                                                                                       
from torch.autograd import Variable                                                                                                                                                                               
from torch.nn.parameter import Parameter                                                                                                                                                                          
                                                                                                                                                                                                                  
                                                                                                                                                                                                                  
class Toy(nn.Module):                                                                                                                                                                                             
    def __init__(self):                                                                                                                                                                                           
        super(Toy, self).__init__()                                                                                                                                                                               
        self.embed = nn.Embedding(100000, 256, sparse=True)                                                                                                                                                       
        self.lin = nn.Linear(256, 256)                                                                                                                                                                            
    def forward(self, idx):                                                                                                                                                                                       
        return self.lin(F.sigmoid(self.embed(idx)))                                                                                                                                                               
                                                                                                                                                                                                                  
                                                                                                                                                                                                                  
toy = Toy().cuda()                                                                                                                                                                                                
optimizer = optim.SparseAdam(toy.parameters(), lr=0.001)                                                                                                                                                          
criterion = nn.CrossEntropyLoss().cuda()                                                                                                                                                                          
                                                                                                                                                                                                                  
                                                                                                                                                                                                                  
x = Variable(torch.from_numpy(numpy.random.randint(0, 100000, 256)).cuda(),                                                                                                                                       
                 requires_grad=False                                                                                                                                                                              
             )                                                                                                                                                                                                    
                                                                                                                                                                                                                  
t = Variable(torch.LongTensor(numpy.random.randint(0, 2, 256)).cuda(),                                                                                                                                            
        requires_grad=False)                                                                                                                                                                                      
                                                                                                                                                                                                                  
                                                                                                                                                                                                                  
start_time = time.time()                                                                                                                                                                                          
for _ in range(2000):                                                                                                                                                                                             
    y = toy(x)                                                                                                                                                                                                    
    cost = criterion(y, t)                                                                                                                                                                                        
    cost.backward()                                                                                                                                                                                               
    optimizer.step()                                                                                                                                                                                              
    print(time.time() - start_time)  

(fails with RuntimeError: SparseAdam does not support dense gradients, please consider Adam instead)