class TemporalPoolerFlockUnit(Unit):
def __init__(self, creator: TensorCreator, params: ExpertParams):
super().__init__(creator.device)
self.flock = TPFlock(params, creator)
def step(self, input_clusters: torch.Tensor, context: torch.Tensor, rewards: torch.Tensor):
self.flock.forward_learn(input_clusters, context, rewards)
def __init__(self, creator: TensorCreator, params: ExpertParams):
super().__init__(creator.device)
self.flock = TPFlock(params, creator)