[Feature Request] Support for Sparse CPU Embeddings in Mixed GPU/CPU model

Thanks for the great work here. It's really exciting to be able to try out mixed-precision so easily now in PyTorch.

Currently I'm training a GPU image classification model (resnet101) but for extreme classification (~3M classes). My model stores the class weights through a CPU based sparse embedding (nn.Embedding(num_proxies, embedding_dim, sparse=True)) where we subsample class weights during training (2048 classes with the positives target weights guaranteed included + some random negative weights) to use for a Sampled Softmax. Each GPU node in DistributedDataParallel will randomly sample class weights independently.

To get this to work, it seems like there's two asks:

1. Can we allow mixed CPU/GPU models to be mixed-precision accelerated? (I think this is just changing the conditions in the code-base from torch.cuda.FloatTensor to also allow torch.FloatTensor).
2. Can we support sparse tensors? (specifically getting RuntimeError: sparse tensors do not have strides from float(model_grad.float().sum()) in scaler.py when using O1 optimization currently)

Example code of model:

class EmbeddingLookup(nn.Module):
    """
    Wrapper for sparse Embedding sampler with custom `embedding_sampler` function
    """
    def __init__(self,
                 num_proxies,
                 embedding_dim,
                 ignore_label=-1.0,
                 embedding_sampler=None):
        super(EmbeddingLookup, self).__init__()

        self.num_proxies = num_proxies
        self.embedding_sampler = embedding_sampler
        self.ignore_label = ignore_label

        self.proxies = nn.Embedding(num_proxies, embedding_dim, sparse=True)

    def forward(self, targets):
        labels, sampled_targets = self.embedding_sampler(targets)
        sampled_proxies = self.proxies(labels)
        return sampled_proxies, sampled_targets

# Add tests once ready to productionize
class ExtremeResNetClassifier(nn.Module):
    def __init__(self,
                 num_classes=0,
                 ignore_label=-1,
                 class_sampler=None,
                 base_model="senet"):

        # enable auto cudnn tuner for fixed image input size
        torch.backends.cudnn.benchmark = True

        super(CPUGPUMixedModel, self).__init__()

        # Only a single GPU per process
        self.device = torch.device('cuda')

        original_model = resnet101(pretrained=pretrained)
        features = nn.Sequential(
            original_model.conv1,
            original_model.bn1,
            original_model.relu,
            original_model.maxpool,
            original_model.layer1,
            original_model.layer2,
            original_model.layer3,
            original_model.layer4,
            nn.AdaptiveAvgPool2d(1)
        )
        distributed = DistributedGPUModel(features).to(self.device)
        self.distributed = DistributedDataParallel(distributed)

        self.class_sample_embeddings = EmbeddingLookup(num_classes, 2048, ignore_label=ignore_label, embedding_sampler=instance_sampler)
        self.softmax_cross_entropy_loss = nn.CrossEntropyLoss(ignore_index=ignore_label).to(self.device)

    def forward(self,
                images,
                class_targets):

        images = images.to(self.device, non_blocking=True)

        embedding = self.distributed(images)
        embedding = embedding.view(embedding.size(0), -1)
        sampled_class_embeddings, sampled_class_targets = self.class_sample_embeddings(class_targets)

        sampled_class_embeddings = sampled_class_embeddings.to(self.device, non_blocking=True)
        sampled_class_targets = sampled_class_targets.to(self.device, non_blocking=True)
        sampled_class_output = nn.functional.linear(embedding, sampled_class_embeddings)
        class_loss = self.softmax_cross_entropy_loss(sampled_class_output, sampled_class_targets)
        return class_loss

Thanks!

Start with scaler.py and the reported O1 failure at float(model_grad.float().sum()), then reproduce the provided mixed GPU/CPU model using a sparse nn.Embedding and sampled softmax. Trace how mixed-device parameters and sparse gradients are handled. Done means the example can train with CPU sparse embeddings and GPU layers under mixed precision without the reported runtime error, with tests added as suggested.