| 128 | self.prev_hidden = None |
| 129 | |
| 130 | def init_hidden(self, batch_size, hidden, height, width, device): |
| 131 | if self.Wci is None: |
| 132 | self.Wci = torch.zeros(1, hidden, height, width, requires_grad=True).to( |
| 133 | device |
| 134 | ) |
| 135 | self.Wcf = torch.zeros(1, hidden, height, width, requires_grad=True).to( |
| 136 | device |
| 137 | ) |
| 138 | self.Wco = torch.zeros(1, hidden, height, width, requires_grad=True).to( |
| 139 | device |
| 140 | ) |
| 141 | return ( |
| 142 | torch.zeros(batch_size, hidden, height, width, requires_grad=True).to( |
| 143 | device |
| 144 | ), |
| 145 | torch.zeros(batch_size, hidden, height, width, requires_grad=True).to( |
| 146 | device |
| 147 | ), |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | class VisionNetwork(nn.Module): |