mirror of
https://github.com/wassname/ray.git
synced 2026-07-12 15:16:46 +08:00
[RLlib] Fix use_lstm flag for ModelV2 (w/o ModelV1 wrapping) and add it for PyTorch. (#8734)
This commit is contained in:
@@ -121,7 +121,7 @@ Once implemented, the model can then be registered and used in place of a built-
|
||||
from ray.rllib.models import ModelCatalog
|
||||
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
|
||||
|
||||
class CustomTorchModel(nn.Module, TorchModelV2):
|
||||
class CustomTorchModel(TorchModelV2):
|
||||
def __init__(self, obs_space, action_space, num_outputs, model_config, name): ...
|
||||
def forward(self, input_dict, state, seq_lens): ...
|
||||
def value_function(self): ...
|
||||
|
||||
Reference in New Issue
Block a user