|
|
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()}) |
|
|
|
|
|
|
