nn.Identity bugfix (#124)

texar/modules/decoders/decoder_base.py

@@ -22,7 +22,7 @@
 import torch
 from torch import nn
 
-from texar.core.layers import identity
+from texar.core.layers import identity, Identity
 from texar.module_base import ModuleBase
 from texar.modules.decoders import decoder_helpers
 from texar.modules.decoders.decoder_helpers import Embedding, Helper
@@ -71,7 +71,7 @@ def _make_output_layer(layer: Optional[Union[nn.Module, torch.Tensor]],
             layer = nn.Parameter(layer, requires_grad=False)
         output_layer.weight = layer
     elif layer is identity:
-        output_layer = identity  # type: ignore
+        output_layer = Identity()
     else:
         raise ValueError(
             f"output_layer should be an instance of `nn.Module`, a tensor,"

texar/modules/encoders/gpt2_encoder.py

@@ -18,7 +18,6 @@
 from typing import Optional
 
 import torch
-import torch.nn as nn
 
 from texar.core import layers
 from texar.hyperparams import HParams
@@ -86,9 +85,9 @@ def __init__(self,
             hparams=self._hparams.position_embed)
 
         # The GPT2 encoder (a TransformerDecoder)
-        self.decoder = TransformerDecoder(
+        self.decoder = TransformerDecoder(  # type: ignore
             vocab_size=self._hparams.vocab_size,
-            output_layer=nn.Identity(),
+            output_layer=layers.identity,
             hparams=self._hparams.decoder)
 
         if self.pretrained_model_dir:

texar/modules/pretrained/gpt2_utils.py

@@ -24,6 +24,7 @@
 import torch
 import torch.nn as nn
 
+from texar.core import layers
 from texar.data.data_utils import maybe_download
 from texar.modules.pretrained.pretrained_utils import default_download_dir
 
@@ -111,7 +112,7 @@ def init_gpt2_checkpoint(model: nn.Module, cache_dir: str):
 
                 output_pointer = name_to_variable(
                     model.decoder, "_output_layer.weight")
-                if not isinstance(output_pointer, nn.Identity):
+                if not isinstance(output_pointer, layers.Identity):
                     assert output_pointer.shape == array.shape
                     output_pointer.data = torch.from_numpy(array)
             elif name == "model/wpe":
@@ -195,7 +196,7 @@ def name_to_variable(model: nn.Module, name: str) -> nn.Module:
             num = int(m_name)
             pointer = pointer[num]  # type: ignore
         else:
-            if not isinstance(pointer, nn.Identity):
+            if not isinstance(pointer, layers.Identity):
                 pointer = getattr(pointer, m_name)
     return pointer