setup.py: Adjustments for ROCm install

requirements/requirements_rocm.txt

@@ -1,3 +1,3 @@
 -r _requirements_base.txt
 --extra-index-url https://download.pytorch.org/whl/rocm6.0
-torch>=2.3.0,<2.4.0
+torch>=2.3.1,<2.4.0

setup.py

@@ -62,8 +62,11 @@
  "zlib-wapi": ("zlib-wapi", ("conda-forge", )),
  "xorg-libxft": ("xorg-libxft", ("conda-forge", ))}
 
-# Include Pillow + Numpy here to avoid installing twice for updated version
-_GROUPS = [["pytorch*", "torch*", "cpuonly", "numpy", "pillow"]]
+# Include Pillow, Numpy and Matplotlib here to avoid installing twice for updated version
+# It goes before pytorch as ROCm torch comes from pip and we don't want the additional packages
+# from pip
+_GROUPS = [["numpy", "pillow", "matplotlib"],
+ ["pytorch*", "torch*", "cpuonly"]]
 """list[list[str]]: Packages that should be installed collectively at the same time """
 
 _DEV_TOOLS = ["flake8", "mypy", "pylint", "pytest", "pytest-mock",
@@ -1089,10 +1092,11 @@ def _install_grouped_packages(self) -> None:
  """ Install packages that should be installed collectively as a group """
  assert self._env.backend is not None
  force_pip = _FORCE_PIP.get("all", []) + _FORCE_PIP.get(self._env.backend, [])
- use_pip = False
- packages = []
- channels: set[str] = set()
  for group in _GROUPS:
+ packages = []
+ channels: set[str] = set()
+ use_pip = False
+
  for item in group:
  for idx, pkg in reversed(list(enumerate(self._packages.to_install))):
  if item != pkg.name and not (item.endswith("*") and
@@ -1105,13 +1109,13 @@ def _install_grouped_packages(self) -> None:
  channels.update(next((v[1] for v in _CONDA_MAPPING.values()
  if v[0] == i_pkg.name),
  ("defaults", )))
- if not packages:
- return
+  if not packages:
+  continue
 
- if use_pip:
- self._from_pip(packages, " ".join(self._packages.pip_arguments))
- else:
- self._from_conda(packages, tuple(channels), conda_only=True)
+  if use_pip:
+  self._from_pip(packages, " ".join(self._packages.pip_arguments))
+  else:
+  self._from_conda(packages, tuple(channels), conda_only=True)
 
  def _install_conda_packages(self) -> None:
  """ Install required conda packages """