akankshakshurai/cuda-error-resolution-analysis huggingface dataset:
Synchronization behavior of torch.cuda.memory_stats and related

username: Zai
cooked: <p>Hi, I have a question about memory_stats and related functions, especially memory_allocated and max_memory_allocated. There seems to be no documentation ( or authoritative information in the forum ) as to whether they require or trigger a synchronization. If not, are their return values just derived from host-available allocator data and information about enqueued ( but potentially not yet executed ) cuda operations ( can they actually give reliable information that way, that won’t possibly change when the operations are actually run )?</p>
[Torch Dynamo] too many subgraphs for LLM model inference due to mask

username: AI79
cooked: <p>Hi,</p>
<p>I plan to use torch.compile on models like Llama  to do some optimizations. For example, to replace the plain implementing of attention to fused attention.<br>
However, dynamo generates hundreds of subgraphs. A lot of them are due to the mask update during each token generation.</p>
<p>Is there anyway to work around it?</p>
<p>Thanks</p>

username: marksaroufim
cooked: <aside class="quote no-group" data-username="AI79" data-post="1" data-topic="185066">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/a4c791/48.png" class="avatar"> AI79:</div>
<blockquote>
<p>I plan to use torch.compile on models like Llama to do some optimizations. For example, to replace the plain implementing of attention to fused attention.<br>
However, dynamo generates hundreds of subgraphs. A lot of them are due to the mask update during each token generation.</p>
<p>Is there anyway to work around it?</p>
<p>Thanks</p>
</blockquote>
</aside>
<p>You can try setting <code>dynamic=True</code>- alternatively i’ve seen some oss implementations preallocate the mask so it’s static</p>
Pytorch2.0 compiled model does not work for huggingface API model.generate()

username: Minerva_Yu
cooked: <p>Hi,<br>
When I’m working with the following code:</p>
<pre><code class="lang-auto">import torch
from transformers import BertTokenizer, BertLMHeadModel
import os
import torch._inductor.config as config

config.debug = True
os.environ["AOT_FX_GRAPHS"] = "1"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', legacy = False)
model = BertLMHeadModel.from_pretrained("bert-base-uncased", torch_dtype=torch.float16).to(device)

text = "test test nonsense words"
encoded_input = tokenizer(text, return_tensors="pt").to(device)

# compile the model
compiled_model = torch.compile(model) 

output_length = 32
input_length = len(encoded_input.input_ids[0])
print("input_ids shape : " + str(encoded_input.input_ids.shape))

# huggingface API generate
generate_ids = compiled_model.generate(encoded_input.input_ids, min_new_tokens =output_length, max_new_tokens=output_length)
output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
print(text)
for outstr in output : 
    print(outstr)
</code></pre>
<p>I want to use the generate() function with the compiled model, but it appears that the compiled model (instance of the class OptimizedModule) does not have the generate() function. Instead, calling generate() triggers _<em>getattr</em> _() and the function is executed by the original BertLMHeadModel.</p>
<p>Is there a method to enable calling generate() with the compiled model?</p>

username: marksaroufim
cooked: <p>Try compiling the generate function directly <code>torch.compile(model.generate)</code></p>
Deploying pytorch 1D CNN model on android app using java

username: dlswja1234
cooked: <p>Hi, I’m trying to deploy 1D CNN based multi-label signal classification model.</p>
<p>However, the output on python code is drastically different with java though I used same model which is deployed using torch.jit.trace.</p>
<p>Please give me any advice I can apply.<br>
I’m trying to solve this problem for one week.</p>

[Python]
<ol>
<li>
<p>model : 1D CNN based DenseNet</p>
</li>
<li>
<p>train data input size : (Batch, channel, length)=(25,1,1000)</p>
</li>
<li>
<p>output class : 4 class / multi label classification</p>
</li>
<li>
<p>version of python and pytorch</p>
</li>
</ol>
<ul>
<li>python : 3.8.8</li>
<li>pytorch : 1.9.0+cu111</li>
</ul>
<ol start="5">
<li>input/output data type : torch_float32</li>
</ol>
<p>[Java]</p>
<ol>
<li>pytorch_android version : 1.9.1</li>
<li>input/output data type : float32(float)</li>
<li>input tensor data :</li>
</ol>
<ul>
<li>real_output : float[1000]</li>
<li>shape : long[]{1,1,1000}</li>
</ul>
<p>[Java Code]<br>
final Tensor inputTensor = Tensor.fromBlob(real_output, shape, MemoryFormat.CONTIGUOUS);<br>
final Tensor ai_output = module.forward(IValue.from(inputTensor)).toTensor();<br>
final float[] ai_result = ai_output.getDataAsFloatArray();<br>
for (int i=0;i&lt;4;i++){<br>
Log.d(TAG, “ai_result_” + i + " : " + ai_result[i]);<br>
}</p>
two sibling layers in RPN are fully connected or 1x1 conv layers?

username: saeed_ni
cooked: <p>I read the faster r-cnn paper and i confused about region proposal section(3.1). now i dont know sibling layer in the RPN are fully connected layers or 1x1 convolutional layers?</p>
<blockquote>
<p>This feature is fed into two sibling fullyconnected layers—a box-regression layer (reg) and a box-classification layer (cls).</p>
</blockquote>
<blockquote>
<p>This architecture is naturally implemented with an n×n convolutional layer followed by two sibling 1 × 1 convolutional layers (for reg and cls, respectively).</p>
</blockquote>
<p>i know project implemented by conv layers. but I cant understand the contradiction in the explanation in the context of the article.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd0f1d65556dd19d1ddac4459d6061d3d8fa0a3f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bd0f1d65556dd19d1ddac4459d6061d3d8fa0a3f" title="121"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd0f1d65556dd19d1ddac4459d6061d3d8fa0a3f_2_472x500.jpeg" alt="121" data-base62-sha1="qYuCLAxy4GWnbAThL9qYZ0GiFST" width="472" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd0f1d65556dd19d1ddac4459d6061d3d8fa0a3f_2_472x500.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd0f1d65556dd19d1ddac4459d6061d3d8fa0a3f.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd0f1d65556dd19d1ddac4459d6061d3d8fa0a3f.jpeg 2x" data-dominant-color="E4E3C1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">121</span><span class="informations">509×539 162 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
Help: cannot pickle 'weakref' object

username: WindyFlyingPig
cooked: <p>I tried to set spawn mode in the beginning of my program by</p>
<pre><code class="lang-auto">    torch.multiprocessing.set_start_method('spawn')
</code></pre>
<p>But when I ran my program, there came an error about ’ cannot pickle ‘weakref’ object’. And it arose at the place where I tried to start multiprocessing function. I thought maybe because the function or args I sent to multiprocessing can’t be pickled. But how can I solve this problem. Thanks!</p>

username: Bhavay_Malhotra
cooked: <p>Hi, I’m facing the same error while my line of code is a bit different.<br>
It is as follows - torch.multiprocessing.set_sharing_strategy(‘file_system’).</p>
<p>What is the possible solution to it?</p>
Hanging on the torch.distributed.init_process_group

username: Flex
cooked: <p>Hi, I am trying to submit slurm job for Pytorch ddp. I use <a href="https://github.com/facebookresearch/XLM/blob/cd281d32612d145c6742b4d3f048f80df8669c30/xlm/slurm.py#L167" rel="noopener nofollow ugc">Facebook Code</a> for configuring distributed training. The program just hangs on the function <code>torch.distributed.init_process_group</code>.The way I submit slurm job is that <code>bash slurm.sh</code> and <code>slurm.sh</code> is</p>
<pre><code class="lang-auto">sbatch -c 16 -n 1 --ntasks-per-node 2 --gres gpu:2 python train.py
</code></pre>
<p>But the slurm job will not hang on the function <code>torch.distributed.init_process_group</code> if I bypass slurm code part of Facebook Code. If I use <code>python -m torch.distributed.launch --nproc_per_node=2 train.py</code> and just call <code>torch.distributed.init_process_group(init_method='env://', backend='nccl')</code> without calling the following functions for submitting slurm jobs:</p>
<pre><code class="lang-auto">hostnames = subprocess.check_output(['scontrol', 'show', 'hostnames', os.environ['SLURM_JOB_NODELIST']])
params.master_addr = hostnames.split()[0].decode('utf-8')
</code></pre>
<p>Does code snippet above give wrong master address so ``torch.distributed.init_process_group` hangs? Code for reproduce is included in <a href="https://github.com/facebookresearch/XLM/blob/cd281d32612d145c6742b4d3f048f80df8669c30/xlm/slurm.py#L167" rel="noopener nofollow ugc">Facebook Code</a>.</p>

username: fduwjj
cooked: <p>Have you tried to verify that all ranks are using the same master_addr?</p>

username: Flex
cooked: <p>I set <code>os.environ['MASTER_ADDR'] = params.master_addr</code>, I guess all ranks have the same master address? Should I print <code>args.master_addr</code> for double checking all ranks?</p>

username: fduwjj
cooked: <p>That’s what I am asking for.</p>
Distribute Dataset Target Transform to GPUs

username: rambo
cooked: <p>Hi!</p>
<p>My <code>Dataset</code> saves a list of parameter values that via some computationally intensive simulation process transform to the high-dimensional tensors that I feed my neural network.</p>
<p>It would be infeasible to save the final tensors in the dataset since I am still working on the simulation process and I would need to regenerate them often before even being able to start training.</p>
<p>I put the generation of the input tensors into the <code>Dataset</code>’s <code>__getitem__</code>-method and accelerate the computation using my GPU and which puts me in the (ill-advised) position of loading my tensors onto the GPU in the <code>Dataset</code> and not after.<br>
In other words, the (target) transform in my <code>Dataset</code> is compute-intensive and happens on the GPU.</p>
<p>I actually have multiple GPUs at my disposal that I would like to take advantage of.<br>
How do I proceed?<br>
I imagine needing to spawn multiple threads, give each their own device (and Dataset? and/or DataLoader?).<br>
How do I combine the outputs of the multiple threads and put them on the device on which the backpropagation happens?</p>
<p>Thank you very much!</p>

username: fduwjj
cooked: <p>Hi, I tried to answer your question but without a detailed context I am not sure if it’s really going to help you.</p>
<p>Can you try to use <code>torch.utils.data.distributed.DistributedSampler</code> in <a href="https://pytorch.org/docs/stable/data.html" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.0 documentation</a> and <code>torch.nn.parallel.DistributedDataParallel</code> in <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel" class="inline-onebox" rel="noopener nofollow ugc">DistributedDataParallel — PyTorch 2.0 documentation</a>?</p>
<p>We need to do multi-process programing here.</p>
Best way to deploy multiple models in one GPU

username: Chen_Li
cooked: <p>I want to deploy 20+ models in one GPU. These models are expected to run in parallel. If the models are deployed in 20+ processes (that’s to say, one model per process), then the GPU complains out of memory (because the initialization step uses a lot of memory). If the models are deployed in one process, then they cannot run in parallel. What’s the best way to do this job?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Chen_Li" data-post="1" data-topic="185382">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/chen_li/48/19110_2.png" class="avatar"> Chen_Li:</div>
<blockquote>
<p>If the models are deployed in one process, then they cannot run in parallel.</p>
</blockquote>
</aside>
<p>Using a single process won’t limit the parallelization of the models and you can still try to use different CUDA streams and launch different models. Depending on the used compute resources the models might be serialized, but this would be the same case for multiple processes.</p>
Multi30k dataset can not load from the torchtext because the main server is down

username: JonathanSum
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/a/2a036f7580b5dddc13aaeda3dc8fe3f6377197d9.png" data-download-href="https://discuss.pytorch.org/uploads/default/2a036f7580b5dddc13aaeda3dc8fe3f6377197d9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a036f7580b5dddc13aaeda3dc8fe3f6377197d9_2_690x274.png" alt="image" data-base62-sha1="5ZFpFzxKsjZsq74O6aL1GplhEmZ" width="690" height="274" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a036f7580b5dddc13aaeda3dc8fe3f6377197d9_2_690x274.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a036f7580b5dddc13aaeda3dc8fe3f6377197d9_2_1035x411.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/a/2a036f7580b5dddc13aaeda3dc8fe3f6377197d9_2_1380x548.png 2x" data-dominant-color="3C3D3D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1797×715 57.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<a href="http://www.quest.dcs.shef.ac.uk" rel="noopener nofollow ugc">www.quest.dcs.shef.ac.uk</a> web site is down.</p>

username: JonathanSum
cooked: <p>Maybe this post is meaningless. But this is not the first time having this. In the last time, it was about a dataset that was hosted on google drive in the TorchText library, and it can not be downloaded at the end.<br>
I hope the server will be recovered very soon.</p>

username: dongqi_cai
cooked: <p>Post a <a href="https://discuss.pytorch.org/t/build-vocab-from-iterator-does-not-work-in-notebook/153575/3">related question link</a> here.<br>
Hope this bug could be fixed soon <img src="https://discuss.pytorch.org/images/emoji/apple/pray.png?v=12" title=":pray:" class="emoji" alt=":pray:" loading="lazy" width="20" height="20"></p>

username: hyejun12123
cooked: <p>I’m having the same problem. The Multi30k function, which worked well until a week ago, is not getting data. Is there a solution?</p>
Gather: gather_list does not contain results from gpu 0

username: thorstenwagner
cooked: <p>I’ve a setup using two GPUs with DistributedDataParallel. I’m using the latest nightly build.</p>
<p>I try to gather the results, but I always endup with 2x the results of gpu 1:</p>
<pre><code class="lang-auto">items = torch.from_numpy(items[:10]).to(self.rank) #necessary because of nvcc, : 10 to make it readable
print(f"Rank {self.rank} items:")
print(items)
if self.rank == 0:
    items_gather_list = [torch.zeros_like(items)] * 2
tdist.barrier()
print("GATHER")
tdist.gather(items,
                    gather_list=items_gather_list if self.rank == 0 else None,
                    dst=0)
if self.rank == 0:
    print(f"LEN GATER LIST", len(items_gather_list))
    print(items_gather_list[0])
    print(" ")
    print(items_gather_list[1])
</code></pre>
<p>That is the output:</p>
<pre><code class="lang-bash">Rank 1 items:                                                                                                                                                                                                                               
Rank 0 items:
tensor([   1, 1073, 1589, 1453, 1997, 1075,  115, 1041, 1133,  756],
       device='cuda:1')
tensor([  44, 1723, 1039, 1321, 1065,  281, 1364, 1830,  296, 1228],
       device='cuda:0')
GATHER
GATHER
LEN GATER LIST 0 2
tensor([   1, 1073, 1589, 1453, 1997, 1075,  115, 1041, 1133,  756],
       device='cuda:0')
 
tensor([   1, 1073, 1589, 1453, 1997, 1075,  115, 1041, 1133,  756],
       device='cuda:0')
</code></pre>
<p>What am I doing wrong?</p>

username: kumpera
cooked: <p>The problem is that you’re passing a list of tensors where all elements are the same tensor.</p>
<pre><code class="lang-auto">items_gather_list = [torch.zeros_like(items)] * 2
</code></pre>
<p>This doesn’t deep copy the tensor, which is what you want.</p>
<p>This should fix it for you:</p>
<pre><code class="lang-auto">items_gather_list = [torch.zeros_like(items), torch.zeros_like(items)] 
</code></pre>

username: thorstenwagner
cooked: <p>Thanks you! It works <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
I have 32 npy files with the shape (32,160,256) and I want to convert them as png images with the same shape

username: Aya_Elfatyany
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/2/02e56677b6a4b29b5e30d93f07e88f1a98baed21.png" data-download-href="https://discuss.pytorch.org/uploads/default/02e56677b6a4b29b5e30d93f07e88f1a98baed21" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/2/02e56677b6a4b29b5e30d93f07e88f1a98baed21_2_690x314.png" alt="image" data-base62-sha1="pCrrTnPfZxbwkg5WP8IquWzsqt" width="690" height="314" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/2/02e56677b6a4b29b5e30d93f07e88f1a98baed21_2_690x314.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/2/02e56677b6a4b29b5e30d93f07e88f1a98baed21_2_1035x471.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/2/02e56677b6a4b29b5e30d93f07e88f1a98baed21_2_1380x628.png 2x" data-dominant-color="1F2020"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2109×962 144 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/f/3fd8d986b7f0d5fe17cc0db8fe5d395d577ecdc7.png" data-download-href="https://discuss.pytorch.org/uploads/default/3fd8d986b7f0d5fe17cc0db8fe5d395d577ecdc7" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/f/3fd8d986b7f0d5fe17cc0db8fe5d395d577ecdc7.png" alt="image" data-base62-sha1="96OG65v6bgeAm3DuluyrzlDW6qz" width="690" height="338" data-dominant-color="272828"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1128×554 42 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: ptrblck
cooked: <p><code>[32, 160, 256]</code> is not a valid image shape so it won’t work. You could create 32 grayscale images from this tensor but not a single image without reducing the channel dimension. It’s also helpful to post code snippets instead of screenshots.</p>
Debugging ONNX Export 'Operator Not Supported' Errors

username: johnorford
cooked: <p>I get an <code>OperatorNotSupported</code> error when trying to do <code>torch.onnx.export</code>.</p>
<p>It was complaining about <code>aten::dict</code>, so I replaced all the dictionary references I could find.</p>
<p>But I still see the same error, and I don’t know where exactly the offending dictionaries are…</p>
<p><code>torch.jit.trace(model,image)</code> doesn’t refer to <code>aten</code> or <code>dict</code> either.</p>

username: johnorford
cooked: <p>OK, I found the <code>verbose</code> flag for the <code>torch.onnx.export</code> function - this should keep me busy for a while!</p>
ValueError: Got dtype 'float64', which is a forbidden dtype

username: alicanakca
cooked: <pre><code class="lang-auto">def __getitem__(self, i: int) -&gt; Tuple[torch.Tensor, torch.Tensor]:
    _id = self.__image_ids[i]
    paths = self.__id2path(_id)
    print("step -1")
    image = self.__load_n_preprocess(paths["image_path"])
    mask = self.__load_n_preprocess(paths["mask_path"])
    print("step 0")
    if self.__transforms is not None:
        segmap = SegmentationMapsOnImage(mask, shape=mask.shape)
        print("step 1")
        image, mask = self.__transforms(image=image, segmentation_maps=segmap)
        print("step 2")
    image = torch.from_numpy(image)
    mask = torch.from_numpy(mask)
    image = image.permute(2, 0, 1)
    mask = mask.permute(2, 0, 1)
    print("step 3")
    return image, mask
</code></pre>
<p>ERROR</p>
<pre><code class="lang-auto">ValueError: Got dtype 'float64', which is a forbidden dtype (bool, uint16, uint32, uint64, uint128, uint256, int32, int64, int128, int256, float16, float32, float64, float96, float128, float256).
</code></pre>
<p>‘data.py’ file includes function above. I detected error with my extraordinary debug style <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> I cannot see “step 2” in results. What should I do ?</p>
<pre><code class="lang-auto">image, mask = self.__transforms(image=image, segmentation_maps=segmap)
</code></pre>
<p>Thanks in advance!</p>

username: bc165870081
cooked: <p>In the imgaug package, for color-related image augmentations, only np.uint8 is supported. Additionally, some other image augmentation methods apart from color-related ones may not necessarily support float64.</p>
<p>I can’t figure out what method you used, but maybe you can check the following URL for confirmation</p>
<p><a href="https://imgaug.readthedocs.io/en/latest/source/dtype_support.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://imgaug.readthedocs.io/en/latest/source/dtype_support.html</a></p>
Distillation of Resnet model (similarity preserving)

username: Johnny_Cash
cooked: <p>Hi, I wanted to do model distillation from the pretrained ResNet model in such a way that the embeddings that were close in the teacher model are also close in the student model. I check the performance of the model afterwards by checking the outputs from the teacher and student after softmax with SRCC against a certain dataset. Somehow my SRCC is around 0 though and I am not sure what I am doing wrong here.</p>
<p>Here is the training loop.</p>
<pre><code class="lang-auto">
teacher_clip = clip_model.visual
student_clip = ModifiedResNet(layers=(2,2,2,2), output_dim=1024, heads=32, input_resolution=224, width=32)

teacher_clip = teacher_clip.to("cuda")
student_clip = student_clip.to("cuda")

max_epochs = 300
ep_log_interval = 1
lrn_rate = 0.005

cel_loss = nn.CosineEmbeddingLoss(reduction="mean")
optimizer = torch.optim.SGD(student_clip.parameters(), lr=lrn_rate)
teacher_clip.eval()

for param in teacher_clip.parameters():
        param.requires_grad = False

for epoch in range(0, max_epochs):
    epoch_loss = 0
    for (batch_idx, batch) in enumerate(loader): 
        with torch.autocast("cuda"):
            X = batch[0].to("cuda")
            Y = teacher_clip(X)

            optimizer.zero_grad()

            oupt = student_clip(X)

            y = torch.ones(Y.shape[0], device=torch.device("cuda:0"))

            loss_val = cel_loss(oupt, Y, y)
            epoch_loss += loss_val.item()

            loss_val.backward()
            optimizer.step()

    if epoch % ep_log_interval == 0:
        print("epoch = %4d   loss = %0.4f" % (epoch, epoch_loss))
        torch.save(student_clip.state_dict(), "./kd_model_1.pt")
</code></pre>