File size: 655 Bytes
23bc32f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from __future__ import annotations
from dataclasses import dataclass
from typing import List
import torch
from .segment import SegmentId
@dataclass
class Batch:
observations: torch.ByteTensor
actions: torch.LongTensor
rewards: torch.FloatTensor
ends: torch.LongTensor
mask_padding: torch.BoolTensor
segment_ids: List[SegmentId]
def pin_memory(self) -> Batch:
return Batch(**{k: v if k == 'segment_ids' else v.pin_memory() for k, v in self.__dict__.items()})
def to(self, device: torch.device) -> Batch:
return Batch(**{k: v if k == 'segment_ids' else v.to(device) for k, v in self.__dict__.items()})
|