Install pytorch3d in aarch64/arm Linux architecture about pytorch3d HOT 4 CLOSED

FROM jupyter/tensorflow-notebook:aarch64-lab-4.0.3
RUN mkdir -p ~/Downloads
RUN cd ~/Downloads
# https://github.com/conda-forge/miniforge/releases/download/24.3.0-0/Mambaforge-24.3.0-0-Linux-aarch64.sh
COPY Mambaforge-24.3.0-0-Linux-aarch64.sh /root/conda.sh
USER root
RUN chmod +x /root/conda.sh && \
    bash /root/conda.sh -b -p /root/conda
RUN /root/conda/bin/conda init
RUN source ~/.bashrc
ENV PATH /root/conda/bin:$PATH
RUN apt install -y vim-tiny
RUN conda create -y -n pytorch3d python=3.10
RUN echo "conda activate pytorch3d" >> ~/.bashrc

RUN CONDA_OVERRIDE_CUDA="12.0" conda install --name pytorch3d --yes pytorch=2.3.1=cuda120_py310h377a36d_200 -c conda-forge -c nvidia

RUN conda install --name pytorch3d --yes -c conda-forge -c fvcore fvcore

RUN /bin/bash -c ". activate pytorch3d && \
    pip install --upgrade pip && \
    pip install \
    scikit-image \
    matplotlib \
    imageio \
    black \
    isort \
    flake8 \
    flake8-bugbear \
    flake8-comprehensions"

# download the pytorch3d v0.7.6 source code from repo
COPY pytorch3d /root/pytorch3d/
# download the torchvision 0.18.1 source code from repo
COPY vision-0.18.1/vision-0.18.1 /root/torchvision/

RUN /bin/bash -c ". activate pytorch3d && \
    pip install ninja && \
    cd /root/torchvision && \
    pip install . "

WORKDIR /root/pytorch3d
# https://developer.nvidia.com/cuda-12-0-0-download-archive
COPY cuda_12.0.0_525.60.13_linux_sbsa.run /root/cuda_12.0.0_525.60.aarch64.run

RUN sh /root/cuda_12.0.0_525.60.aarch64.run --toolkit  --silent
ENV CUDA_HOME /usr/local/cuda-12.0

RUN /bin/bash -c ". activate pytorch3d && \
    cd /root/pytorch3d && \
    FORCE_CUDA=1 pip install . "