opacus does not support the torch.nn.Embedding module about opacus HOT 4 CLOSED

import torch
import numpy as np

class FeaturesLinear(torch.nn.Module):
    def __init__(self, field_dims, output_dim=1):
        super().__init__()
        self.fc = torch.nn.Embedding(sum(field_dims), output_dim)
        self.bias = torch.nn.Parameter(torch.zeros((output_dim, )))
        if np.__version__ < "1.24":
            self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)
        else:
            self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.int64)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = x + x.new_tensor(self.offsets).unsqueeze(0)
        return torch.sum(self.fc(x), dim=1) + self.bias


class Model(torch.nn.Module):
    """
    A pytorch implementation of Logistic Regression.
    """
    def __init__(self, **kwargs):
        super().__init__()
        embedding_nums = 100000
        if 'embedding_nums' in kwargs:
            embedding_nums = int(kwargs['embedding_nums'])
        input_dim = 10
        if 'input_dim' in kwargs:
            input_dim = int(kwargs['input_dim'])
        self.linear = FeaturesLinear(field_dims=[embedding_nums for _ in range(input_dim)])

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        return torch.sigmoid(self.linear(x).squeeze(1))