RuntimeError: don't know how to restore data location of torch.FloatStorage

I am trying to train Yolov5 using DML and after tweaks in the code to let it use the DML device, I have this error when it tries to load the weights on the GPU:

 File "D:\Documenti\models\GTA_code\dml\lib\site-packages\torch\serialization.py", line 592, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "D:\Documenti\models\GTA_code\dml\lib\site-packages\torch\serialization.py", line 851, in _load
    result = unpickler.load()
  File "D:\Documenti\models\GTA_code\dml\lib\site-packages\torch\serialization.py", line 843, in persistent_load
    load_tensor(data_type, size, key, _maybe_decode_ascii(location))
  File "D:\Documenti\models\GTA_code\dml\lib\site-packages\torch\serialization.py", line 832, in load_tensor
    loaded_storages[key] = restore_location(storage, location)
  File "D:\Documenti\models\GTA_code\dml\lib\site-packages\torch\serialization.py", line 812, in restore_location
    return default_restore_location(storage, str(map_location))
  File "D:\Documenti\models\GTA_code\dml\lib\site-packages\torch\serialization.py", line 178, in default_restore_location
    raise RuntimeError("don't know how to restore data location of "
RuntimeError: don't know how to restore data location of torch.FloatStorage (tagged with dml)

Loading pretrained weights is not a supported feature…?

[smk2007]

Hi, and thanks for posting the issue.

We have not taken a look at yolov5 yet, and as such the operators needed to get yolov5 working end-to-end are spotty. We absolutely intend to support yolov5 in the future, and will update this thread when we expect to have the operator conformance needed to enable it.

[lyleaf-81]

Hello, do you support YOLOV4?