Function
update_dataset
(layer, dataset, dev, attention_mask, position_ids)
Source from the content-addressed store, hash-verified
| 20 | import os |
| 21 | |
| 22 | def update_dataset(layer, dataset, dev, attention_mask, position_ids): |
| 23 | with torch.no_grad(): |
| 24 | with torch.cuda.amp.autocast(): |
| 25 | for index, inps in enumerate(dataset): |
| 26 | inps = inps.to(dev) |
| 27 | if len(inps.shape)==2: |
| 28 | inps = inps.unsqueeze(0) |
| 29 | new_data = layer(inps, attention_mask=attention_mask,position_ids=position_ids)[0].to('cpu') |
| 30 | dataset.update_data(index,new_data) |
| 31 | |
| 32 | |
| 33 | def block_ap( |
Tested by
no test coverage detected