How to save the optimized texture map?

[winterwindwang]

I used the test code from the #710, as shown as follows
assset test.zip

import kaolin
from PIL import Image
import torch
import numpy as np
from matplotlib import pyplot as plt
from torchvision import transforms, models

import kaolin as kal

mesh_path = r"new_musk.obj"
mesh = kaolin.io.obj.import_mesh(mesh_path, with_materials=True, with_normals=True)

vertices = mesh.vertices.cuda().unsqueeze(0)
vertices.requires_grad = False
faces = mesh.faces.cuda()
uvs = mesh.uvs.cuda().unsqueeze(0)

nb_faces = faces.shape[0]
# nb_vertices = vertices_init.shape[1]
face_size = 3

face_uvs_idx = mesh.face_uvs_idx.cuda()


def get_camera_from_view(elev, azim, r=3.0, look_at_height=0.0):
    x = r * torch.sin(elev) * torch.sin(azim)
    y = r * torch.cos(elev)
    z = r * torch.sin(elev) * torch.cos(azim)

    pos = torch.tensor([x, y, z]).unsqueeze(0)
    look_at = torch.zeros_like(pos)
    look_at[:, 1] = look_at_height
    direction = torch.tensor([0.0, 1.0, 0.0]).unsqueeze(0)

    camera_proj = kal.render.camera.generate_transformation_matrix(pos, look_at, direction)
    return camera_proj


face_uvs = kal.ops.mesh.index_vertices_by_faces(uvs, face_uvs_idx).detach()
face_uvs.requires_grad = False
texture_map = mesh.materials[0]['map_Kd'].unsqueeze(0).float().permute(0, 3, 1, 2).cuda() / 255
texture_map.requires_grad = True

texture_map = torch.autograd.Variable(texture_map.cuda(), requires_grad=True)
optimizer = torch.optim.Adam([texture_map], lr=0.01)

criterion = torch.nn.CrossEntropyLoss()

model = models.resnet50(pretrained=True)
model.eval()
model.cuda()

for idx in range(100):
    # camera projection
    camera = kal.render.camera.generate_perspective_projection(np.pi / 3).cuda()
    camera_projection = camera
    elev = 1.047
    azim = 0.0
    radius = 1.5
    look_at_height = 0.25

    camera_transform = get_camera_from_view(torch.tensor(elev), torch.tensor(azim), r=radius,
                                            look_at_height=look_at_height).cuda()
    face_vertices_camera, face_vertices_image, face_normals = kal.render.mesh.prepare_vertices(
        vertices, faces, camera_projection, camera_transform=camera_transform)
    ### Perform Rasterization ###
    # Construct attributes that DIB-R rasterizer will interpolate.
    # the first is the UVS associated to each face
    # the second will make a hard segmentation mask
    face_attributes = [
        face_uvs.repeat(1, 1, 1, 1),
        torch.ones((1, nb_faces, 3, 1), device='cuda')
    ]

    image_features, soft_mask, face_idx = kal.render.mesh.dibr_rasterization(
        224, 224, face_vertices_camera[:, :, :, -1],
        face_vertices_image, face_attributes, face_normals[:, :, -1],
        rast_backend='cuda')
    # image_features is a tuple in composed of the interpolated attributes of face_attributes

    texture_coords, mask = image_features

    mask = mask.cpu()
    image = kal.render.mesh.texture_mapping(texture_coords, texture_map.repeat(1, 1, 1, 1),
                                            mode='bilinear')
    image = torch.clamp(image * mask.cuda(), 0., 1.)

    image = image.permute((0, 3, 1, 2))
    output = model(image)

    loss = criterion(output, torch.tensor([200]).cuda())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(f"Iter: {idx}, Loss: {loss.item()}")
transforms.ToPILImage()(texture_map[0]).save("optimized_texture.png")

Although the texture of the rendered image has been modified, the saved optimized texture map is the same as the original one. Could you tell me how to save the optimized texture map correctly? Thank you!