| 75 | """ |
| 76 | |
| 77 | def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): |
| 78 | factory_kwargs = {"dtype": dtype, "device": device} |
| 79 | super().__init__() |
| 80 | self.linear_1 = nn.Linear( |
| 81 | in_features=in_channels, |
| 82 | out_features=hidden_size, |
| 83 | bias=True, |
| 84 | **factory_kwargs, |
| 85 | ) |
| 86 | self.act_1 = act_layer() |
| 87 | self.linear_2 = nn.Linear( |
| 88 | in_features=hidden_size, |
| 89 | out_features=hidden_size, |
| 90 | bias=True, |
| 91 | **factory_kwargs, |
| 92 | ) |
| 93 | |
| 94 | def forward(self, caption): |
| 95 | hidden_states = self.linear_1(caption) |