Reinforcement Learning Unity ML Agents

TypeError: invalid dtype object: only floating-point types are supported as the default type
I changed in ml-agents/ml-agents/mlagents/torch_utils/torch.py
if _device.type == “cuda”:
torch.set_default_device(_device.type)
torch.set_default_dtype(torch.float32) // changed
else:
torch.set_default_dtype(torch.float32)
logger.debug(f"default Torch device: {_device}")
After changing I get this error:
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

Seems setting threaded: false in the config solves the problem

behaviors:
  SnowballTarget:
    trainer_type: ppo
    summary_freq: 10000
    keep_checkpoints: 10
    checkpoint_interval: 50000
    max_steps: 200000
    time_horizon: 64
    threaded: false
    hyperparameters:
      learning_rate: 0.0003
      learning_rate_schedule: linear
      batch_size: 128
      buffer_size: 2048
      beta: 0.005
      epsilon: 0.2
      lambd: 0.95
      num_epoch: 3
    network_settings:
      normalize: false
      hidden_units: 256
      num_layers: 2
      vis_encode_type: simple
    reward_signals:
      extrinsic:
        gamma: 0.99
        strength: 1.0