[BUG]When I use the am model to run a PCTSP experiment, an error is reported：TypeError("_DeviceDtypeModuleMixin.__init__() got an unexpected keyword argument 'use_native_sdpa'")

[lihaoya5]

Describe the bug

Here's the command I used and the error message：
E:\rl4co-main>python run.py experiment=archive/pctsp/am

Traceback (most recent call last):
File "C:\Users\qian.conda\envs\rl6co\Lib\site-packages\hydra_internal\instantiate_instantiate2.py", line 92, in _call_target
return target(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\rl4co-main\rl4co\models\zoo\am\model.py", line 33, in init
super().init(env, policy, baseline, baseline_kwargs, **kwargs)
File "E:\rl4co-main\rl4co\models\rl\reinforce\reinforce.py", line 40, in init
super().init(env, policy, **kwargs)
File "E:\rl4co-main\rl4co\models\rl\common\base.py", line 73, in init
super().init(**litmodule_kwargs)
File "C:\Users\qian.conda\envs\rl6co\Lib\site-packages\lightning\pytorch\core\module.py", line 123, in init
super().init(*args, **kwargs)
TypeError: _DeviceDtypeModuleMixin.init() got an unexpected keyword argument 'use_native_sdpa'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "E:\rl4co-main\rl4co\utils\utils.py", line 36, in wrap
metric_dict, object_dict = task_func(cfg=cfg)
^^^^^^^^^^^^^^^^^^
File "E:\rl4co-main\rl4co\tasks\train.py", line 44, in run
model: LightningModule = hydra.utils.instantiate(cfg.model, env)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\qian.conda\envs\rl6co\Lib\site-packages\hydra_internal\instantiate_instantiate2.py", line 226, in instantiate
return instantiate_node(
^^^^^^^^^^^^^^^^^
File "C:\Users\qian.conda\envs\rl6co\Lib\site-packages\hydra_internal\instantiate_instantiate2.py", line 347, in instantiate_node
return _call_target(target, partial, args, kwargs, full_key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\qian.conda\envs\rl6co\Lib\site-packages\hydra_internal\instantiate_instantiate2.py", line 97, in _call_target
raise InstantiationException(msg) from e
hydra.errors.InstantiationException: Error in call to target 'rl4co.models.zoo.am.model.AttentionModel':
TypeError("_DeviceDtypeModuleMixin.init() got an unexpected keyword argument 'use_native_sdpa'")
full_key: model
[2024-03-11 16:49:09,246][rl4co.utils.utils][INFO] - Output dir: E:\rl4co-main\logs\train\runs\pctsp20\am-pctsp20\2024-03-11_16-49-06
[2024-03-11 16:49:09,246][rl4co.utils.utils][INFO] - Closing loggers...
Error executing job with overrides: ['experiment=archive/pctsp/am']
Error in call to target 'rl4co.models.zoo.am.model.AttentionModel':
TypeError("_DeviceDtypeModuleMixin.init() got an unexpected keyword argument 'use_native_sdpa'")
full_key: model

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

my python environment:

[Federico Berto]

Problem solved, please follow the new instructions here! 👍