| 55 | # |
| 56 | |
| 57 | class MyModule(nn.Module): |
| 58 | def __init__(self, in_features: int, out_features: int, bias: bool = True): |
| 59 | super(MyModule, self).__init__() |
| 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 | ###################################################################### |