Include clean PRIMA code and evaluator; document metrics, initialization, and MIT contributions
91c1e9e verified | """Move M4SEG batches without changing mask or index dtypes.""" | |
| import torch | |
| def move_batch(batch, device): | |
| def move(value): | |
| if torch.is_tensor(value): | |
| return value.to(device=device, non_blocking=True) | |
| if isinstance(value, dict): | |
| return {key: move(item) for key, item in value.items()} | |
| if isinstance(value, list): | |
| return [move(item) for item in value] | |
| if isinstance(value, tuple): | |
| return tuple(move(item) for item in value) | |
| return value | |
| result = move(batch) | |
| for key in ('global_enc_images', 'grounding_enc_images'): | |
| if result[key] is not None: | |
| result[key] = result[key].to(dtype=torch.bfloat16) | |
| return result | |