#include "RandomNetwork.hpp"

RandomNetworkImpl::RandomNetworkImpl(std::map<std::string,std::size_t> nbOutputsPerState) : nbOutputsPerState(nbOutputsPerState)
{
}

torch::Tensor RandomNetworkImpl::forward(torch::Tensor input)
{
  if (input.dim() == 1)
    input = input.unsqueeze(0);

  return torch::randn({input.size(0), (long)nbOutputsPerState[getState()]}, torch::TensorOptions().device(device).requires_grad(true));
}

std::vector<std::vector<long>> RandomNetworkImpl::extractContext(Config &, Dict &) const
{
  return std::vector<std::vector<long>>{{0}};
}