What OS are you using? There may be issue with NVIDIA CUDA driver

nvidia-smi verifies that NVIDIA driver is installed. To verify CUDA driver , you can use binaries from CUDA samples

Also, did you install NVIDIA container toolkit and configure your Docker engine?

Example

% docker run --rm --gpus all nvcr.io/nvidia/k8s/cuda-sample:devicequery-cuda12.5.0
Unable to find image 'nvcr.io/nvidia/k8s/cuda-sample:devicequery-cuda12.5.0' locally
devicequery-cuda12.5.0: Pulling from nvidia/k8s/cuda-sample
a8b1c5f80c2d: Pull complete 
5b5f281aa9fa: Pull complete 
e85ef0eeb66c: Pull complete 
0d0a907de515: Pull complete 
508725ab4fe4: Pull complete 
4f4fb700ef54: Pull complete 
9420ad0cd1fa: Pull complete 
02166da33929: Pull complete 
012a2ab2c8b7: Pull complete 
968b9e56c870: Pull complete 
Digest: sha256:ac53daee629763d712e1361b77e4c4f4ad146148f9dffc6288a75732270c6e85
Status: Downloaded newer image for nvcr.io/nvidia/k8s/cuda-sample:devicequery-cuda12.5.0
/cuda-samples/sample Starting...

 CUDA Device Query (Runtime API) version (CUDART static linking)

Detected 1 CUDA Capable device(s)

Device 0: "Tesla T4"
  CUDA Driver Version / Runtime Version          12.6 / 12.5
  CUDA Capability Major/Minor version number:    7.5
  Total amount of global memory:                 14916 MBytes (15640231936 bytes)
  (040) Multiprocessors, (064) CUDA Cores/MP:    2560 CUDA Cores
  GPU Max Clock rate:                            1590 MHz (1.59 GHz)
  Memory Clock rate:                             5001 Mhz
  Memory Bus Width:                              256-bit
  L2 Cache Size:                                 4194304 bytes
  Maximum Texture Dimension Size (x,y,z)         1D=(131072), 2D=(131072, 65536), 3D=(16384, 16384, 16384)
  Maximum Layered 1D Texture Size, (num) layers  1D=(32768), 2048 layers
  Maximum Layered 2D Texture Size, (num) layers  2D=(32768, 32768), 2048 layers
  Total amount of constant memory:               65536 bytes
  Total amount of shared memory per block:       49152 bytes
  Total shared memory per multiprocessor:        65536 bytes
  Total number of registers available per block: 65536
  Warp size:                                     32
  Maximum number of threads per multiprocessor:  1024
  Maximum number of threads per block:           1024
  Max dimension size of a thread block (x,y,z): (1024, 1024, 64)
  Max dimension size of a grid size    (x,y,z): (2147483647, 65535, 65535)
  Maximum memory pitch:                          2147483647 bytes
  Texture alignment:                             512 bytes
  Concurrent copy and kernel execution:          Yes with 3 copy engine(s)
  Run time limit on kernels:                     No
  Integrated GPU sharing Host Memory:            No
  Support host page-locked memory mapping:       Yes
  Alignment requirement for Surfaces:            Yes
  Device has ECC support:                        Enabled
  Device supports Unified Addressing (UVA):      Yes
  Device supports Managed Memory:                Yes
  Device supports Compute Preemption:            Yes
  Supports Cooperative Kernel Launch:            Yes
  Supports MultiDevice Co-op Kernel Launch:      Yes
  Device PCI Domain ID / Bus ID / location ID:   0 / 0 / 30
  Compute Mode:
     < Default (multiple host threads can use ::cudaSetDevice() with device simultaneously) >

deviceQuery, CUDA Driver = CUDART, CUDA Driver Version = 12.6, CUDA Runtime Version = 12.5, NumDevs = 1
Result = PASS

Note the Result = PASS output.

Above was executed on AL2023 on g4dn. Driver and NVIDIA toolkit installed as per this article

Thank you for providing detailed information about your GPU instance and the issue you're experiencing with CUDA driver initialization in Docker containers. Based on the information you've shared, it appears that while your host system is correctly recognizing the GPU, there's a problem with the CUDA environment inside the Docker containers. Here are some steps and suggestions to help resolve this issue:

1. NVIDIA Container Toolkit: Ensure that you have the latest version of the NVIDIA Container Toolkit installed. This toolkit is crucial for GPU support in Docker containers. You can update it using the following commands:

   sudo dnf config-manager --add-repo https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo
   sudo dnf install -y nvidia-container-toolkit
   

2. Docker Runtime Configuration: After installing or updating the NVIDIA Container Toolkit, make sure to configure the Docker runtime:

   sudo nvidia-ctk runtime configure --runtime=docker
   sudo systemctl restart docker
   

3. Docker Run Command: When running your Docker container, ensure you're using the --gpus all flag correctly. Your command looks good, but double-check that it's being used consistently.

4. CUDA Version Compatibility: The PyTorch image you're using (pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime) is built for CUDA 12.1, while your host system has CUDA 12.6. This version mismatch could potentially cause issues. Try using a PyTorch image that matches your host CUDA version more closely.

5. Environment Variables: Sometimes, setting specific environment variables can help. Try adding these to your Docker run command:

   -e NVIDIA_VISIBLE_DEVICES=all -e NVIDIA_DRIVER_CAPABILITIES=all
   

6. Host Driver Version: Your host driver version (560.35.03) is recent, which is good. However, ensure it's compatible with both your host CUDA version and the CUDA version in your Docker image.

7. Container Toolkit Configuration: Verify that the NVIDIA Container Toolkit is correctly configured by running:

   sudo nvidia-ctk runtime configure --runtime=docker
   

8. Test with NVIDIA CUDA Container: As a diagnostic step, try running an official NVIDIA CUDA container to isolate whether the issue is specific to PyTorch or more general:

   docker run --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi
   

If these steps don't resolve the issue, you may need to investigate further by checking Docker and NVIDIA Container Runtime logs. Also, ensure that your instance's Amazon ECS-optimized AMI is up-to-date, as newer versions include improvements for GPU support.

If the problem persists after trying these suggestions, please provide the output of the following commands from both the host and inside a Docker container:

• nvidia-smi
• nvcc -V
• Docker logs when attempting to run a GPU-enabled container

This additional information will help in further diagnosing the issue.
Sources
How do I install NVIDIA GPU driver, CUDA toolkit, NVIDIA Container Toolkit on Amazon EC2 instances running Amazon Linux 2023 (AL2023)? | AWS re:Post
NVIDIA drivers for your Amazon EC2 instance - Amazon Elastic Compute Cloud
New ECS-Optimized AMIs support AL2 Kernel 5.10 for GPU and Inf Instances