(self, input, mask)
| 60 | self.linear = nn.Linear(in_features, out_features, bias) |
| 61 | |
| 62 | def forward(self, input, mask): |
| 63 | with profiler.record_function("LINEAR PASS"): |
| 64 | out = self.linear(input) |
| 65 | |
| 66 | with profiler.record_function("MASK INDICES"): |
| 67 | threshold = out.sum(axis=1).mean().item() |
| 68 | hi_idx = np.argwhere(mask.cpu().numpy() > threshold) |
| 69 | hi_idx = torch.from_numpy(hi_idx).cuda() |
| 70 | |
| 71 | return out, hi_idx |
| 72 | |
| 73 | |
| 74 | ###################################################################### |
nothing calls this directly
no outgoing calls
no test coverage detected