GPU-to-NPU Migration: Handwritten Digit Recognition

Scenario

In the context of training and inference for AI foundation models, Ascend NPUs, a Chinese-developed computing platform, deliver performance comparable to NVIDIA GPUs. However, developers using existing PyTorch code built for GPUs often worry about the high migration costs and the complexity of adaptation when switching to this platform. The torch_npu package, an adaptation layer for Ascend NPUs, allows standard PyTorch code to run on NPU hardware with minimal modifications, significantly reducing the migration barrier. This tutorial is designed for developers who need to port existing GPU-based PyTorch training code to the Ascend NPU environment.

This tutorial uses MNIST handwritten digit recognition as an example to demonstrate how to migrate PyTorch training code from GPUs to Ascend NPUs. By highlighting the differences between GPU and NPU implementations, this tutorial helps developers quickly understand the migration process and complete training and inference verification on the ModelArts platform.

Concepts

Migration Principles

torch_npu works by creating an adaptation layer between the PyTorch framework and the Ascend CANN software stack. This layer automatically translates PyTorch's CUDA calls into equivalent operations on the Ascend NPU. As a result, developers can migrate models without rewriting the core logic; they only need to replace device-specific API calls to complete the migration.

Typical migration scenario: If your service needs to migrate from an NVIDIA GPU cluster to an Ascend NPU cluster, such as for localization or cost optimization, you can use this guide to quickly adapt your existing PyTorch training and inference code to the NPU environment.

Prerequisites

• You have registered a Huawei Cloud account, completed real-name authentication, and been granted related permissions. For details, see 1. Prerequisites: Accounts and Permissions.
• The notebook instance is in the Running state and meets the following requirements. For details about how to create a notebook instance, see Creating a Notebook Instance.

  The following table describes the instance specifications, images, and disk specifications.

  Parameter	NPU Environment	GPU Environment	Description
  Instance Specifications	You are advised to use a single PU (Ascend: 1*ascend-snt9b1) or a higher configuration.	You are advised to use a single-node single-PU flavor or a higher configuration.	The GPU is optional and is used for cross-validation.
  Image	Select the preset ModelArts image pytorch_2.7.1-cann_8.3.rc1-py_3.11-hce_2.0.2509-aarch64-snt9b.	Select the preset ModelArts image pytorch_1.12.1-cuda_10.2-py_3.9.11-ubuntu_20.04-x86_64.	Select an image of the corresponding architecture.
  Disk Specifications	5 GB or larger is recommended.	5 GB or larger is recommended.	The disk is used to store datasets and model weights.

Billing

Running notebook instances in ModelArts development environments use compute and storage resources, which are billed. For details about the billing, see Development Environment Billing Item.

Precautions

Before the migration, back up the original GPU code and model weight file. Architecture differences (precaution 1) will make x86 images completely unavailable on NPUs. Before creating an instance, ensure that the image architecture matches the target architecture to avoid resource waste.

Core Concepts

The core of this case is a closed-loop process: "environment verification → code migration → training verification → inference comparison." This process provides a clear and intuitive understanding of the entire process of migrating PyTorch code from GPUs to NPUs.

Core of migration: Only three modifications are required (see Step 3).

• 1. Add import torch_npu after import torch.
• 2. Replace the device identifier torch.cuda with torch.npu.
• 3. Replace the device identifier cuda with npu. The code for defining the network model, training logic, and evaluation logic does not need to be modified.

Universal weight: The model weight file (.pt) saved by PyTorch is universal between GPUs and NPUs. You only need to adjust the map_location parameter to load the file across devices.

Monitoring comparison: Use npu-smi info for NPUs to replace nvidia-smi of GPUs. AICore(%) corresponds to the NPU usage, GPU-Util corresponds to the GPU usage, and HBM-Usage(MB) corresponds to the memory usage.

Step 1: Access a Notebook Instance

1. Log in to the ModelArts console and perform the following operations as required:
   • New console: Choose Model Build > Notebook.
   • Old console: Choose Development Space > Notebook.
2. In the Operation column of the target notebook instance, click Access Environment. In the Access Method dialog box, click Access on the right of JupyterLab Access in the WebIDE tab.
3. Go to the JupyterLab page. The ModelArts Launcher page is automatically displayed. In the Notebook area of the ModelArts Launcher page, click PyTorch to create a PyTorch notebook instance.
   Figure 1 Opening PyTorch
   

Step 2: Verify the Environment

Before starting the migration, verify that the current environment is available. After confirming that the device type and computing capability are normal, proceed with the subsequent code migration.

1. Check the device.
   Run the following commands to automatically check whether the current environment uses GPUs, NPUs, or CPUs:

   import torch
   
   if torch.cuda.is_available():
       device = torch.device("cuda")
       print(f"Current environment: GPU - {torch.cuda.get_device_name(0)}")
   elif hasattr(torch, 'npu') and torch.npu.is_available():
       import torch_npu
       device = torch.device("npu")
       print(f"Current environment: NPU - {torch.npu.get_device_name(0)}")
   else:
       device = torch.device("cpu")
       print("Current environment: CPU")
   
   print(f"Device identifier: {device}")

   After the code is executed, the current device type and device identifier are displayed. The following is an example. If "Current environment: CPU" is displayed, check whether the image is correctly selected or whether the resource specifications are successfully allocated.

   • Output example - NPU:

     Current environment: NPU - xxx
     Device identifier: npu

   • Output example - GPU:

     Current environment: GPU - xxx
     Device identifier: cuda

2. Verify the matrix operation.

   Verify the device computing capability through matrix multiplication. The following examples show GPU and NPU environments, respectively. Execute them based on the current operating environment.

   • GPU matrix operation verification:

     import torch
     try:
         _has_gpu = torch.cuda.is_available()
     except Exception:
         _has_gpu = False
     if _has_gpu:
         print("===== GPU matrix operation =====")
         x = torch.randn(10000, 10000).cuda()
         y = torch.randn(10000, 10000).cuda()
         z = x.mm(y)
         print(f"Result shape: {z.shape}")
          print(f"Device location: {z.device}")
          print("GPU matrix operation verification passed!")
     else:
         print("No GPU is available in the current environment. GPU matrix operation verification is skipped.")

     Expected output:

     ===== GPU matrix operation =====
     Result shape: torch.Size([10000, 10000])
     Device location: cuda:0
     GPU matrix operation verification passed!

   • NPU matrix operation verification:

     import torch
     try:
         import torch_npu
         _has_npu = torch.npu.is_available()
     except Exception:
         _has_npu = False
     if _has_npu:
         print("===== NPU matrix operation =====")
         x = torch.randn(10000, 10000).npu()
         y = torch.randn(10000, 10000).npu()
         z = x.mm(y)
         print(f"Result shape: {z.shape}")
          print(f"Device location: {z.device}")
         print("NPU matrix operation verification passed!")
     else:
         print("No NPU is available in the current environment. NPU matrix operation verification is skipped.")

     Expected output:

     ===== NPU matrix operation =====
     Result shape: torch.Size([10000, 10000])
     Device location: npu:0
     NPU matrix operation verification passed!

   Key points in migration: In the NPU environment, you only need to add import torch_npu and replace .cuda () with .npu (). The rest of the code remains unchanged.

   Verification method: If the device is available, the message "Matrix operation verification passed!" and the result shape torch.Size([10000, 10000]) should be displayed.

Step 3: Perform MNIST Handwritten Digit Recognition

This step completes the entire training process for MNIST handwritten digit recognition, including defining the network model, preparing and loading data, defining training and evaluation functions, and performing GPU/NPU-based training. The only difference lies in the training code, which varies between GPUs and NPUs. All other code remains consistent across both environments.

1. Define a network model.

   The definition of the network model is hardware-agnostic, meaning the abstract definition remains consistent across different devices (GPUs or NPUs). This case employs a LeNet-5 variant convolutional neural network designed to classify 28 x 28 grayscale images. The architecture comprises two convolutional layers (Conv2d), two Dropout layers, and two fully connected layers (Linear), ultimately producing a probability distribution across 10 classes.

   import torch
   import torch.nn as nn
   import torch.nn.functional as F
   
   class Net(nn.Module):
       def __init__(self):
           super(Net, self).__init__()
           self.conv1 = nn.Conv2d(1, 32, 3, 1)
           self.conv2 = nn.Conv2d(32, 64, 3, 1)
           self.dropout1 = nn.Dropout(0.25)
           self.dropout2 = nn.Dropout(0.5)
           self.fc1 = nn.Linear(9216, 128)
           self.fc2 = nn.Linear(128, 10)
   
       def forward(self, x):
           x = self.conv1(x)
           x = F.relu(x)
           x = self.conv2(x)
           x = F.relu(x)
           x = F.max_pool2d(x, 2)
           x = self.dropout1(x)
           x = torch.flatten(x, 1)
           x = self.fc1(x)
           x = F.relu(x)
           x = self.dropout2(x)
           x = self.fc2(x)
           output = F.log_softmax(x, dim=1)
           return output
   
   print("Network model defined.")

   Expected output:

   Network model defined.

2. Prepare data.
   Download the MNIST dataset and build a data loader. The code in the GPU and NPU environments is the same. The MNIST dataset can be downloaded in either of the following ways. Select a method based on your network environment.

   Download Method	Application Scenario	Description
   Method 1: Directly download the dataset.	The network can access GitHub.	Download the raw data file from GitHub.
   Method 2: Download the dataset from an OBS bucket.	The network is restricted.	Download the packaged data from the Huawei Cloud OBS bucket.

   • Method 1: Directly download the dataset.

     # Standard mode
     !mkdir -p ./data/MNIST/raw
     !wget https://raw.githubusercontent.com/leo-mao/MNIST_data/master/train-images-idx3-ubyte.gz -O ./data/MNIST/raw/train-images-idx3-ubyte.gz
     !wget https://raw.githubusercontent.com/leo-mao/MNIST_data/master/train-labels-idx1-ubyte.gz -O ./data/MNIST/raw/train-labels-idx1-ubyte.gz
     !wget https://raw.githubusercontent.com/leo-mao/MNIST_data/master/t10k-images-idx3-ubyte.gz -O ./data/MNIST/raw/t10k-images-idx3-ubyte.gz
     !wget https://raw.githubusercontent.com/leo-mao/MNIST_data/master/t10k-labels-idx1-ubyte.gz -O ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz

     import os
     import gzip
     from torchvision import datasets, transforms
     os.makedirs('./data/MNIST/raw', exist_ok=True)
     
     # Decompress the .gz file.
     for f in os.listdir('./data/MNIST/raw/'):
         if f.endswith('.gz'):
             with gzip.open(f'./data/MNIST/raw/{f}', 'rb') as gz:
                 with open(f'./data/MNIST/raw/{f[:-3]}', 'wb') as out:
                     out.write(gz.read())

   • Method 2: Download the dataset from an OBS bucket.

     # Download the dataset from the OBS bucket.
     !pip install download
     
     %env no_proxy='a.test.com,127.0.0.1,2.2.2.2'
     from download import download
     # Download the MNIST dataset.
     url = "https://eduhicomputing.obs.cn-south-1.myhuaweicloud.com:443/%20ComputerVision2/MNIST_Data.zip"
     path = download(url, "./data/MNIST/raw", kind="zip", replace=True)
     !mv ./data/MNIST/raw/MNIST_Data/test/* ./data/MNIST/raw/MNIST_Data/train/* ./data/MNIST/raw/ && rm -rf ./data/MNIST/raw/MNIST_Data

3. Load the dataset. Use torchvision to load the MNIST dataset and create the DataLoader for the training and test sets.

   from torchvision import datasets, transforms
   transform = transforms.Compose([
       transforms.ToTensor(),
       transforms.Normalize((0.1307,), (0.3081,))
   ])
   dataset1 = datasets.MNIST('./data', train=True, download=False, transform=transform)
   dataset2 = datasets.MNIST('./data', train=False, download=False, transform=transform)
   print(f"Training set size: {len(dataset1)}")
   print(f"Test set size: {len(dataset2)}")

   Expected output:

   Training set size: 60000
   Test set size: 10000

   If the data fails to be loaded, check whether the data file is completely downloaded.

4. Define training and evaluation functions. Define the training function train and evaluation function test. The functions automatically adapt to GPUs or NPUs based on the device parameter in the request. The function logic does not need to be modified.

   The training function moves data and tags to the specified device in each batch, and the evaluation function calculates the loss and accuracy. The two functions use the device parameter to adapt to different devices and do not need to be modified.

   def train(model, device, train_loader, optimizer, epoch):
       """Train the model and return the average loss and accuracy."""
       model.train()
       total_loss = 0
       correct = 0
       total = 0
       for batch_idx, (data, target) in enumerate(train_loader):
           data, target = data.to(device), target.to(device)
           optimizer.zero_grad()
           output = model(data)
           loss = F.nll_loss(output, target)
           loss.backward()
           optimizer.step()
           total_loss += loss.item()
           pred = output.argmax(dim=1, keepdim=True)
           correct += pred.eq(target.view_as(pred)).sum().item()
           total += len(data)
       avg_loss = total_loss / len(train_loader)
       accuracy = 100. * correct / total
       print(f'Train Epoch: {epoch} | Loss: {avg_loss:.4f} | Accuracy: {correct}/{total} ({accuracy:.0f}%)')
       return avg_loss, accuracy
   
   
   def test(model, device, test_loader):
       """Evaluate the model and return the average loss and accuracy."""
       model.eval()
       test_loss = 0
       correct = 0
       with torch.no_grad():
           for data, target in test_loader:
               data, target = data.to(device), target.to(device)
               output = model(data)
               test_loss += F.nll_loss(output, target, reduction='sum').item()
               pred = output.argmax(dim=1, keepdim=True)
               correct += pred.eq(target.view_as(pred)).sum().item()
   
       test_loss /= len(test_loader.dataset)
       accuracy = 100. * correct / len(test_loader.dataset)
       print(f'Test  | Average loss: {test_loss:.4f} | Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.0f}%)')
       return test_loss, accuracy
   
   print("Training and evaluation functions defined.")

   Expected output:

   Training and evaluation functions defined.

5. Perform GPU/NPU-based training. Use GPUs and NPUs to train the model and compare the code before and after the migration. Save the model weight file after training for three epochs.
   • GPU-based training: The following code runs on a GPU-based instance.

     import torch
     import torch.optim as optim
     from torch.optim.lr_scheduler import StepLR
     try:
         import torch.cuda 
         cuda_available = torch.cuda.is_available()
     except (ImportError, AttributeError):
         cuda_available = False
     
     # ====== GPU-based training ======
     if cuda_available:
         device = torch.device("cuda")
         print(f"Device used: {device} - {torch.cuda.get_device_name(0)}")
         train_kwargs = {'batch_size': 64, 'shuffle': True}
         test_kwargs = {'batch_size': 1000, 'shuffle': False}
         train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
         test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
         model = Net().to(device)
         optimizer = optim.Adadelta(model.parameters(), lr=1.0)
         scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
         for epoch in range(1, 4):
             train(model, device, train_loader, optimizer, epoch)
             test(model, device, test_loader)
             scheduler.step()
         torch.save(model.state_dict(), "mnist_gpu.pt")
         print("GPU-based training completed. The model has been saved as mnist_gpu.pt.")
     else:
         print("No GPU is available in the current environment. GPU-based training is skipped.")

     Expected output:

     Devices used: cuda - Tesla V100S-PCIE-32GB
     Train Epoch: 1 | Loss: 0.1968 | Accuracy: 56419/60000 (94%)
     Test  | Average loss: 0.0542 | Accuracy: 9826/10000 (98%)
     Train Epoch: 2 | Loss: 0.0759 | Accuracy: 58647/60000 (98%)
     Test  | Average loss: 0.0417 | Accuracy: 9867/10000 (99%)
     Train Epoch: 3 | Loss: 0.0580 | Accuracy: 59008/60000 (98%)
     Test  | Average loss: 0.0313 | Accuracy: 9891/10000 (99%)
     GPU-based training completed. The model has been saved as mnist_gpu.pt.

     Verification method: If GPU image resources are used, the message "GPU-based training completed. The model has been saved as mnist_gpu.pt." should be displayed after training. The test accuracy for all three epochs should be within the range of 97% to 99%. If the message "No GPU available in the current environment." is displayed, check whether the instance specifications include GPUs.

   • NPU-based training: The following code runs on an NPU-based instance.

     Key points for migration: Compared with GPU-based training, only three changes are made:

     1. Add import torch_npu.
     2. torch.cuda→torch.npu
     3. cuda→npu

     import torch
     import torch.optim as optim
     from torch.optim.lr_scheduler import StepLR
     try:
         import torch_npu # [Change 1] Import the NPU adaptation package.
         npu_available = torch.npu.is_available() # [Change 2] torch.cuda → torch.npu
     except (ImportError, AttributeError):
         npu_available = False
     
     # ====== NPU-based training ======
     if npu_available:
         device = torch.device("npu")   # [Change 3] cuda → npu
         print(f"Device used: {device} - {torch.npu.get_device_name(0)}")
     
         train_kwargs = {'batch_size': 64, 'shuffle': True}
         test_kwargs = {'batch_size': 1000, 'shuffle': False}
     
         train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
         test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
     
         model = Net().to(device)
         optimizer = optim.Adadelta(model.parameters(), lr=1.0)
         scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
     
         for epoch in range(1, 4):
             train(model, device, train_loader, optimizer, epoch)
             test(model, device, test_loader)
             scheduler.step()
     
         torch.save(model.state_dict(), "mnist_npu.pt")
         print("NPU-based training completed. The model has been saved as mnist_npu.pt.")
     else:
         print("No NPU is available in the current environment. NPU-based training is skipped.")

     Expected output:

     Device used: npu - Ascend910B4
     Train Epoch: 1 | Loss: 0.2067 | Accuracy: 56290/60000 (94%)
     Test  | Average loss: 0.0442 | Accuracy: 9844/10000 (98%)
     Train Epoch: 2 | Loss: 0.0761 | Accuracy: 58734/60000 (98%)
     Test  | Average loss: 0.0344 | Accuracy: 9885/10000 (99%)
     Train Epoch: 3 | Loss: 0.0544 | Accuracy: 59053/60000 (98%)
     Test  | Average loss: 0.0319 | Accuracy: 9890/10000 (99%)
     NPU-based training completed. The model has been saved as mnist_npu.pt.

     Verification method: If NPU image specifications are used, the message "NPU-based training completed. The model has been saved as mnist_npu.pt." should be displayed after training. The test accuracy for all three epochs should be basically the same as that for GPU-based training. If the message "No NPU available in the current environment." is displayed, check whether the CANN version is selected for the image.

Step 4: Perform Model Inference and Cross-Device Loading

In this step, the trained model is used for inference and verification, including loading the model weight file, predicting images in the test set, and visualizing the results. The PyTorch weight file is compatible between GPUs and NPUs. You only need to adjust the map_location parameter to load the file across devices.

1. Load the trained model weight file for inference. The code automatically detects the current device type and loads the corresponding model weight.

   The weight file (.pt) saved by PyTorch is compatible between GPUs and NPUs. You only need to adjust the map_location parameter.

   Key points for migration: The model weight file in this case is universal. In the NPU scenario, only map_location needs to be changed from cuda:0 to npu:0. In addition, the code for network definition, training logic, evaluation logic, and data loading is the same.

   import torch
   
   # Automatically detect the device and load the corresponding model.
   if torch.cuda.is_available():
       device = torch.device("cuda")
       model_path = "mnist_gpu.pt"
       map_location = "cuda:0"  
   elif hasattr(torch, 'npu') and torch.npu.is_available():
       import torch_npu
       device = torch.device("npu")
       model_path = "mnist_npu.pt"
       map_location = "npu:0" # In the NPU scenario, change map_location from "cuda:0" to "npu:0".
   else:
       device = torch.device("cpu")
       model_path = "mnist_gpu.pt"
       map_location = "cpu"
   
   model = Net().to(device)
   model.load_state_dict(torch.load(model_path, map_location=map_location))
   model.eval()
   print(f"Model loaded. Device: {device}")

   Expected output:

   • NPU:

     Model loaded. Device: npu

   • GPU:

     Model loaded. Device: cuda

2. Select 12 images from the test dataset, use the loaded model to perform inference, and display the comparison between the prediction results and the actual tags.

   import matplotlib.pyplot as plt
   num = 12
   cols = 4
   rows = (num + cols - 1) // cols
   fig, axes = plt.subplots(rows, cols, figsize=(15, 6))
   model.eval()
   with torch.no_grad():
       for i in range(num):
           image, label = dataset2[i]
           output = model(image.unsqueeze(0).to(device))
           pred = output.argmax(dim=1).item()
           ax = axes[i // cols, i % cols]
           ax.imshow(image.squeeze(), cmap='gray')
           color = 'green' if pred == label else 'red'
           ax.set_title(f'L:{label} P:{pred}', color=color, fontsize=10)
           ax.axis('off')
   plt.tight_layout()
   plt.show()

   Expected output:

   

   In the prediction result graph, the green title indicates a correct prediction, and the red title indicates an incorrect prediction.

Step 5: Monitor the Device

During training and inference, you can run device monitoring commands to check the compute usage and memory usage of NPUs or GPUs. This helps you determine whether the training is running properly. For GPUs, run the nvidia-smi command. For NPUs, run the npu-smi info command. The following table lists the core monitoring metrics of the two types of devices.

• GPU monitoring

  Run the nvidia-smi command on the GPU instance to check the GPU usage and memory usage.

  import torch
  
  if torch.cuda.is_available():
      !nvidia-smi
  else:
      print("No GPU is available in the current environment. nvidia-smi is skipped.")

  Expected output:

  Wed Jul 29 11:50:13 2026       
  +-----------------------------------------------------------------------------+
  | NVIDIA-SMI 470.57.02    Driver Version: 470.57.02    CUDA Version: 11.4     |
  |-------------------------------+----------------------+----------------------+
  | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
  | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
  |                               |                      |               MIG M. |
  |===============================+======================+======================|
  |   0  xxx...  On   | 00000000:00:0D.0 Off |                    0 |
  | N/A   35C    P0    37W / 250W |   2925MiB / 32510MiB |      0%      Default |
  |                               |                      |                  N/A |
  +-------------------------------+----------------------+----------------------+
  
  +-----------------------------------------------------------------------------+
  | Processes:                                                                  |
  |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
  |        ID   ID                                                   Usage      |
  |=============================================================================|
  +-----------------------------------------------------------------------------+

• NPU monitoring

  Run the npu-smi info command on the NPU instance to check the NPU usage and memory usage. AICore(%) indicates the compute usage of NPUs, which is equivalent to GPU-Util of GPUs. HBM-Usage(MB) indicates the memory usage, which is equivalent to Memory-Usage of GPUs.

  import torch
  
  if hasattr(torch, 'npu') and torch.npu.is_available():
      !npu-smi info
  else:
      print("No NPU is available in the current environment. npu-smi info is skipped.")

  Expected output:

  +------------------------------------------------------------------------------------------------+
  | npu-smi 25.5.1                   Version: 25.5.1                                               |
  +---------------------------+---------------+----------------------------------------------------+
  | NPU   Name                | Health        | Power(W)    Temp(C)           Hugepages-Usage(page)|
  | Chip                      | Bus-Id        | AICore(%)   Memory-Usage(MB)  HBM-Usage(MB)        |
  +===========================+===============+====================================================+
  | 0     xxx               | OK            | 87.5        42                0    / 0             |
  | 0                         | 0000:C1:00.0  | 0           0    / 0          4822 / 32768         |
  +===========================+===============+====================================================+
  +---------------------------+---------------+----------------------------------------------------+
  | NPU     Chip              | Process id    | Process name             | Process memory(MB)      |
  +===========================+===============+====================================================+
  | 0       0                 | 3180          | python                   | 1999                    |
  +===========================+===============+====================================================+

FAQ

• What do I do if a message is displayed indicating that no NPU is available in the current environment when I run the NPU training code?

  Check the following items in sequence:

  1. Whether the NPU type is selected for the resource specifications of the notebook instance.
  2. Whether the AArch64 image that contains CANN is selected.
  3. Whether the instance has been started and NPU resources have been allocated.
• Can the model weights saved during GPU-based training be directly loaded on NPUs?

  Yes. The PyTorch weight file is device-agnostic. You only need to adjust the map_location parameter. In some scenarios, you need to load the weights into CPUs by setting map_location=cpu and then convert them to NPUs.

Helpful Links