[RLlib] Implement DQN PyTorch distributional head. (#9589)

2026-09-09 11:32:43 +08:00 · 2020-07-25 09:29:24 +02:00
parent 9f3570828a
commit 5dc4b6686e
12 changed files with 358 additions and 187 deletions
@@ -206,6 +206,10 @@ def get_variable(value,
    elif framework == "torch" and torch_tensor is True:
        torch, _ = try_import_torch()
        var_ = torch.from_numpy(value)
+        if dtype == torch.float32:
+            var_ = var_.float()
+        elif dtype == torch.int32:
+            var_ = var_.int()
        if device:
            var_ = var_.to(device)
        var_.requires_grad = trainable