After training a word embedding model on a large-ish corpus, my embeddings converge to nan values.
The model is very simple (skipgram with negative sampling)
full model:
class NEG_loss(nn.Module):
    def __init__(self, vocab_size, embed_size, neg_sampling_table=None):
        """
        :param vocab_size: An int. The number of possible classes.
        :param embed_size: An int. EmbeddingLockup size
        :param num_sampled: An int. The number of sampled from noise examples
        :param neg_sampling_table: A list of non negative floats. Class neg_sampling_table. None if
            using uniform sampling. The neg_sampling_table are calculated prior to
            estimation and can be of any form, e.g equation (5) in [1]
        """

        super(NEG_loss, self).__init__()

        self.device = t.device("cuda:0" if t.cuda.is_available() else "cpu")

        self.vocab_size = vocab_size
        self.embed_size = embed_size

        self.out_embed.weight = nn.Parameter(
            t.cat(
                [
                    t.zeros(1, self.embed_size),
                    FT(self.vocab_size - 1, self.embed_size).uniform_(
                        -0.5 / self.embed_size, 0.5 / self.embed_size
                    ),
                ]
            )
        )


        self.in_embed.weight = nn.Parameter(
            t.cat(
                [
                    t.zeros(1, self.embed_size),
                    FT(self.vocab_size - 1, self.embed_size).uniform_(
                        -0.5 / self.embed_size, 0.5 / self.embed_size
                    ),
                ]
            )
        )

        self.neg_sampling_table = neg_sampling_table
        if self.neg_sampling_table is not None:
            assert min(self.neg_sampling_table) >= 0, "Each weight should be >= 0"

            self.neg_sampling_table = Variable(t.from_numpy(neg_sampling_table)).float()

    # TODO this is bad - find more elegant solution
    def sample(self, num_sample):
        """
        draws a sample from classes based on neg_sampling_table
        """

        return self.neg_sampling_table[
            t.randint(0, len(self.neg_sampling_table), (num_sample,))
        ]

    def forward(self, input_labels, out_labels, num_sampled):
        """
        :param input_labels: Tensor with shape of [batch_size] of Long type
        :param out_labels: Tensor with shape of [batch_size, window_size] of Long type
        :param num_sampled: An int. The number of sampled from noise examples
        :return: Loss estimation with shape of [1]
            loss defined in Mikolov et al. Distributed Representations of Words and Phrases and their Compositionality
            papers.nips.cc/paper/5021-distributed-representations-of-words-and-phrases-and-their-compositionality.pdf
        """

        input_labels = input_labels.to("cuda:0")
        out_labels = out_labels.to("cuda:0")

        batch_size = out_labels.size()[0]

 
        input_ = self.in_embed(input_labels.view(-1))
        output = self.out_embed(out_labels.view(-1))

        if self.neg_sampling_table is not None:
           
            noise_sample_count = batch_size * num_sampled
            draw = self.sample(noise_sample_count)
            noise = draw.view(batch_size, num_sampled).long()
            
        else:


            noise = Variable(
                t.Tensor(batch_size, num_sampled)
                .uniform_(0, self.vocab_size - 1)
                .long()
            )
        # if use_cuda:
        noise = noise.to(self.device)
        noise = self.out_embed(noise).neg()

        log_target = (input_ * output).sum(1).squeeze().sigmoid().log()

        """ ∑[batch_size * window_size, num_sampled, embed_size] * [batch_size * window_size, embed_size, 1] ->
            ∑[batch_size, num_sampled, 1] -> [batch_size] """
        sum_log_sampled = (
            t.bmm(noise, input_.unsqueeze(2)).sigmoid().log().sum(1).squeeze()
        )

        loss = log_target + sum_log_sampled

        return -loss.mean()

    def input_embeddings(self):
        return self.in_embed.weight.detach().cpu().numpy()

and training loop and optimizer code
    # NEG loss and optim
    neg = NEG_loss(vocab_size, dim, neg_sampling_table=neg_sampling_dist)
    neg.to("cuda:0")
    optimizer = Adam(neg.parameters(), 0.01)

    sys.stdout.write("BEGINNING TRAINING\n")
    for i in range(epochs):
        sys.stdout.write("-" * 35 + "\n")
        sys.stdout.write(f"EPOCH {i+1}\n")
        for i, batch in enumerate(tqdm(dataloader)):
            input_, output_ = batch
            optimizer.zero_grad()
            loss = neg(input_, output_, neg_samples)
            loss.backward()
            optimizer.step()
        sys.stdout.write("-" * 35 + "\n")

This model is about as simple as it gets so I am a bit surprised I am having this issue. Everything is fine on a smaller corpus (7k batches of size 128) but gets nan-y with a larger corpus (200k batches of size 128).
Anyone see anything immediately wrong or have any tips to figure out what is going?
Any help is appreciated. Also pretty new to pytorch - so if you see anything that is a bit dumb feel free to point it out