akankshakshurai/cuda-error-resolution-analysis huggingface dataset:
Torch ppc64le wheels

WillTrojak: <p>Hi,</p>
<p>I have access to IBM power9 and power10 hardware, however there aren’t prebuilt wheels for that hardware so I’ve been compiling it from source. I was wondering if I could contribute some wheels for that hardware?</p>
<p>At the moment you have wheels such as</p>
<pre><code class="lang-auto">torch-2.4.0-cp311-cp311-manylinux1_x86_64.whl
torch-2.4.0-cp311-cp311-manylinux2014_aarch64.whl
torch-2.4.0-cp311-cp311-win_amd64.whl
torch-2.4.0-cp311-none-macosx_11_0_arm64.whl
</code></pre>

rc1092: <p>HI ,<br>
I have been trying to install pytorch from source without any success. I am on power9 (RHEL 8). would you be able to please share the wheels. \</p>
<p>Thank you</p>
Matmul slows down when doing communication overlapping

phucnh: <blockquote>
<p>The classical example is a NCCL communication kernel running in the background. In such cases the matmul expects to be able to use all the SMs but is prevented from doing so because some of the are busy. This can lead to its blocks being scheduled as two separate waves on the available SMs. This “wave quantization” can double the latency of the matmul kernels.</p>
</blockquote>
<p>Hi, I wonder if there is a workaround for the issue described here <a href="https://github.com/pytorch/pytorch/pull/144974" rel="noopener nofollow ugc">[link]</a> when doing comm/computation overlapping (TLDR: when overlapping communication and computation, matmul becomes slower) ?</p>

fduwjj: <p>Maybe you want to allocate dedicated SMs for communication and limit the SM usage of communication so that you don’t see slowing in matmul that much?</p>
Introduction to Libuv TCPStore Backend

githubsgi: <p>Thanks for the <a href="https://github.com/pytorch/tutorials/blob/main/intermediate_source/TCPStore_libuv_backend.rst" rel="noopener nofollow ugc">article</a>. Wondering if you can provide some details about the content of the TCPStore and what is its role in c10d .</p>
<p>N.B. Previously posted at <a href="https://github.com/pytorch/tutorials/issues/3272" class="inline-onebox" rel="noopener nofollow ugc">Introduction to Libuv TCPStore Backend · Issue #3272 · pytorch/tutorials · GitHub</a> .</p>

fduwjj: <p>I see, yeah agree we might need a generic tutorial for tcpstore as well. cc: <a class="mention" href="/u/d4l3k">@d4l3k</a></p>
I want to collbrate with you

Ujjwal_Yadav: <p>Dear PYTORCH  Partnerships Team,</p>
<p>Greetings from ISIEINDIA (Imperial Society of Innovative Engineers – India)!</p>
<p>We are India’s leading organization dedicated to nurturing innovation, skill development, and entrepreneurship among youth through experiential learning. With over 1,200+ institutional collaborations, 75,000+ trained students, and support from national bodies like AICTE, MSME, and the Ministry of Education, ISIEINDIA has been at the forefront of bridging academia and industry.</p>
<p>We are reaching out with a proposal to collaborate with pytorch as an Education Partner for our new initiative — AI Labs — which we are deploying in schools and universities across India. These labs are designed to train students in AI, robotics, digital tools, and entrepreneurship.</p>
<p>A key part of this program includes teaching students how to visually present ideas, design brand identities, pitch decks, and marketing content. We see pytorch as an essential tool in this journey.</p>
<p>This partnership would give pytoch visibility and usage across thousands of educational institutions while empowering the next generation of creators, designers, and entrepreneurs in India.</p>
<p>Please let us know the appropriate contact or next steps to explore this collaboration further. We’re happy to setup a detailed meeting at your convenience.</p>
<p>Looking forward to hearing from you.</p>
<p>ujjwal kumar singh</p>
<p>ISIEINDIA</p>
<p>Phone: +91 9773678725</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://isieindia.com/">
  <header class="source">
      <img src="https://isieindia.com/wp-content/uploads/fbrfg/favicon-32x32.png" class="site-icon" width="32" height="32">

      <a href="https://isieindia.com/" target="_blank" rel="noopener nofollow ugc">Best EV Training Platform</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:600/338;"><img src="https://isieindia.com/wp-content/uploads/2024/11/New-Pages-Update-43.jpg" class="thumbnail" width="85" height="85"></div>

<h3><a href="https://isieindia.com/" target="_blank" rel="noopener nofollow ugc">Front Page</a></h3>

  <p>India's 1st Job Leading EV Courses in India These Programs are certified by ASDC (Automotive Skill Development Council) and Endorsed by SMEV.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>ReplyReply AllForwardEdit as new</p>

ptrblck: <p>CC <a class="mention" href="/u/matt_white">@Matt_White</a> for visibility</p>

Matt_White: <p><a class="mention" href="/u/ujjwal_yadav">@Ujjwal_Yadav</a> Could you please email <a href="mailto:training@pytorch.org">training@pytorch.org</a>? thank you.</p>
Multiple forwards and comp graph building

iliasslasri: <p>Hello everyone,<br>
I have one question. Does multiple similar forwards with different inputs/outputs trigger multiple builds of the computational graph knowing that the graph is exactly the same, if yes, how much this affect the memory usage?</p>
<p>Thanks!</p>

ptrblck: <p>Yes, different forward passes will create different computation graphs, which will store the different intermediate forward activations (which are needed for the gradient computation in the backward pass). The increase in device memory depends on the model architecture and the size of forward activations which need to be stored.</p>
<p>Here is a simple example using a <code>ResNet50</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torchvision.models as models 

print(torch.cuda.memory_allocated() / 1024**2)
# 0.0

device = "cuda"
model = models.resnet50().to(device)
x = torch.randn(1, 3, 224, 224, device=device)

print(torch.cuda.memory_allocated() / 1024**2)
#98.30224609375

out1 = model(x)
print(torch.cuda.memory_allocated() / 1024**2)
# 189.58642578125

out2 = model(x)
print(torch.cuda.memory_allocated() / 1024**2)
# 271.33935546875
</code></pre>
<p>Calling <code>backward</code> will free the intermediate forward activations, but the first <code>backward</code> call will also create the <code>grad</code> attributes which will use memory.</p>

iliasslasri: <p>Agreed, now I’m certain! Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
Torch.fx symbolic_trace failing with TypeError on torch.ones: “slice indices must be integers or None”

SchrodingersCat: <p>Hi folks! I am getting going with PyTorch. I’m building a tool that wraps HuggingFace models in a custom <code>WrappedModel</code> so I can trace their execution using <code>torch.fx.symbolic_trace</code>. The goal is to analyze the traced graph and detect certain ops like <code>float32</code> usage.</p>
<p>To do this, I:</p>
<ul>
<li>Wrap the model in a subclass of <code>torch.nn.Module</code>.</li>
<li>Run a <code>forward()</code> pass with dummy <code>input_ids</code>.</li>
<li>Call <code>symbolic_trace(wrapped_model)</code> or fall back to <code>torch.jit.trace()</code>.</li>
</ul>
<p><strong>What’s going wrong:</strong></p>
<p>I consistently see:</p>
<blockquote>
<p>Forward pass failed in WrappedModel — slice indices must be integers or None or have an <strong>index</strong> method</p>
</blockquote>
<p>And ultimately:</p>
<blockquote>
<p>Rule run failed: ‘method’ object is not iterable</p>
</blockquote>
<p><strong>Likely problematic code:</strong></p>
<pre><code class="lang-auto">class WrappedModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_ids):
        try:
            batch_size = input_ids.size(0)
            seq_len = input_ids.size(1)

            # This line fails during symbolic tracing
            attention_mask = torch.ones((batch_size, seq_len), dtype=torch.int64)

            output = self.model(input_ids=input_ids, attention_mask=attention_mask)
        except Exception as e:
            logging.warning(f"TRACE ERROR inside wrapped forward: {e}")
            return torch.zeros(1, 1)

        if hasattr(output, "last_hidden_state"):
            return output.last_hidden_state
        elif hasattr(output, "logits"):
            return output.logits
        return output

</code></pre>
<p><strong>What I have already tried:</strong></p>
<ul>
<li>Using <code>input_ids.size(0)</code> instead of <code>input_ids.shape[0]</code></li>
<li>Making sure the dummy input has fixed dimensions: <code>torch.randint(0, 1000, (1, 10))</code></li>
<li>Hardcoding the mask shape (e.g., <code>torch.ones((1, 10), dtype=torch.int64)</code>)</li>
<li>Falling back to <code>torch.jit.trace</code> — same error during forward</li>
<li>Switching between <code>BertModel</code> and <code>BertForSequenceClassification</code></li>
</ul>
<p><strong>What (I think) I am asking for:</strong></p>
<p>How do I make <code>torch.ones(...)</code> work inside a traced wrapper model during <code>symbolic_trace()</code>?</p>
<p>Thank you in advance for any guidance.</p>
Is there an out-of-place equivalent to boolean indexing assignment?

fouronnes: <p>Is there an out-of-place equivalent to boolean indexing assignment <code>A[M] = B</code>?</p>
<p>For example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

A = torch.arange(10)
M = A % 2 == 0
B = torch.full_like(A, 10)[M]

# Can I do this out of place?
A[M] = B

# is there a function like C = foo(A, M, B)
# note that A and B don't have the same shape, so I can't use torch.where
</code></pre>

fouronnes: <p>Answering my own question a few months later because I figured it out. There is indeed two options:</p>
<ul>
<li><a href="https://pytorch.org/docs/stable/generated/torch.Tensor.index_put.html" rel="noopener nofollow ugc">torch.index_put</a></li>
<li><a href="https://pytorch.org/docs/stable/generated/torch.scatter.html" rel="noopener nofollow ugc">torch.scatter</a></li>
</ul>
<p>Example with <code>torch.index_put()</code>:</p>
<pre><code class="lang-auto">A = torch.arange(10)
M = A % 2 == 0
B = torch.full_like(A, 10)[M]

indices = M.nonzero().squeeze(-1)

C = A.index_put((indices,), B)
</code></pre>

fouronnes: <p>Or even better with <code>masked_scatter</code>:</p>
<pre><code class="lang-auto">A = torch.arange(10)
M = A % 2 == 0
B = torch.full_like(A, 10)[M]

C = A.masked_scatter(M, B)
</code></pre>
Discrepancy Between Theoretical and Measured FLOPs/token for LLaMA-4 Scout 17B (MoE)

Arunima_Ghosh: <p>Hi everyone,</p>
<p>I’m currently working with the meta-llama/Llama-4-Scout-17B-16E model for text generation and analyzing its performance. To estimate theoretical compute requirements, I’ve been using the following simplified formula for FLOPs/token:</p>
<p>Estimated Simplified FLOPs/token (in GFLOPs) = 2N + 2 * L * S * H</p>
<p>Where:</p>
<ul>
<li><strong>N</strong> = Total model parameters (in billions)</li>
<li><strong>L</strong> = Number of transformer layers</li>
<li><strong>S</strong> = Sequence length (input + output)</li>
<li><strong>H</strong> = Hidden size</li>
</ul>
<p>However, I’m noticing a <strong>significant discrepancy</strong> between this estimated value and the <strong>measured average FLOPs/token</strong> calculated as:</p>
<p>Measured FLOPs/token = Total_FLOPs / (batch_size * (input_len + output_len))<br>
Here, I have calculated Total_FLOPS using pytorch profiler.</p>
<p>Even accounting for sparsity due to Mixture-of-Experts (MoE) (e.g., 1 out of 16 experts active), the <strong>measured FLOPs/token is much lower</strong> than the simplified theoretical value. This is especially true when running inference with standard batch sizes and sequence lengths.</p>
<p>I’m using LLaMA-4 purely for <strong>text generation</strong>, so the <strong>visual transformer</strong> part of the configuration (present in the model’s config) shouldn’t be active. I’d like to confirm:</p>
<ul>
<li>Is the simplified formula inherently overestimating for MoE models?</li>
<li>Does the sparsity in MoE layers (1/16 experts) reduce compute more drastically than reflected in the formula?</li>
<li>Are there internal optimizations (e.g., kv-caching, fused ops) that significantly reduce measured FLOPs/token?</li>
</ul>
<p>Any insights or corrections to the formula (specific to MoE-based decoder-only models like LLaMA-4) would be greatly appreciated.</p>
<p>Thanks</p>
New PyTorch Sphinx Theme - Feedback Requested

svekars: <p>We’re excited to announce that the updated PyTorch Sphinx theme is now available in the <a href="https://pytorch.org/docs/main/" rel="noopener nofollow ugc">main</a> branch and will be available in the stable branch starting version 2.8! The new theme is based on the pydata_sphinx_theme and includes the following new features:</p>
<ul>
<li>Dark mode</li>
<li>Page rating capability</li>
<li>Expandable left navigation</li>
<li>Feedback button linking to GitHub issue creation</li>
<li>Edit on GitHub button</li>
<li>Horizontal navigation bar</li>
</ul>
<h2><a name="p-469561-feedback-needed-1" class="anchor" href="#p-469561-feedback-needed-1"></a>Feedback Needed</h2>
<p>Please try out the theme and fill out <a href="https://forms.gle/VJCypjGdZ1TyJkuE6" rel="noopener nofollow ugc">our survey</a> with your feedback. You can also submit an issue directly to pytorch/pytorch with the label: <strong>sphinx-theme-feedback</strong>.</p>
<p>Thank you!</p>
How to efficiently tune Hyperparameters?

Abas1: <p>I’m fine-tuning EfficientNet-B0 on an imbalanced dataset (5 classes, 73% majority class) with 35K total images. Currently using 10% of data for faster iteration.</p>
<p>I’m balancing various hyperparameters and extras :</p>
<ul>
<li>Learning rate</li>
<li>Layer unfreezing schedule</li>
<li>Learning rate decay rate/timing</li>
<li>optimizer</li>
<li>different pretrained models(not a hyperparameter)</li>
</ul>
<p>How can I systematically understand the impact of each hyperparameter without explosion of experiments? Is there a standard approach to isolate parameter effects while maintaining computational efficiency?</p>
<p>Currently I’m changing one parameter at a time (e.g., learning decay rate from 0.1→0.3) and running short training runs, but I’d appreciate advice on best practices. How do you prevent the scenario of making multiple changes and running full 60-epoch training only to not know which change was responsible for improvements? Would it be better to first run a baseline model on the full dataset for 50+ epochs to establish performance, then identify which hyperparameters most need optimization, and only then experiment with those specific parameters on a smaller subset?</p>
<p>How do people train for 1000 Epochs confidently?</p>
<pre><code class="lang-auto"></code></pre>
Looping over DataLoader Returns List?

Imahn: <p>Hi,</p>
<p>I have a PyTorch tensor and use <code>torch.utils.data.TensorDataset</code> to make it a <code>torch.utils.data.Dataset</code> object. Then I throw it into a <code>DataLoader</code> and loop over it, yet the data I obtain is a <em>list</em>, but a tensor, and I don’t get yet why. Here an MWE:</p>
<pre><code class="lang-auto">import torch 

data_geant_eval = torch.ones(256)
test_loader = torch.utils.data.DataLoader(
    dataset=torch.utils.data.TensorDataset(data_geant_eval),
    batch_size=256, 
)
for (idx, data) in enumerate(test_loader): 
    print(type(data))
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">&lt;class 'list'&gt;
</code></pre>
<p>Is that expected behavior?</p>

ptrblck: <p>Yes, this is expected since <code>TensorDataset.__getitem__</code> returns a <code>tuple</code> as seen <a href="https://github.com/pytorch/pytorch/blob/28ed0bdb3701cdf71b2e281a4c00f7318db8c98a/torch/utils/data/dataset.py#L195-L196">here</a>.</p>

Imahn: <p>Thank you for the answer! This is the <code>_getitem__</code> function of <code>TensorDataset</code>:</p>
<pre><code class="lang-auto">def __getitem__(self, index):
        return tuple(tensor[index] for tensor in self.tensors)
</code></pre>
<p>It’s supposed to return a tuple, so I’m a bit confused on why my MWE returns the type of a list.</p>

ptrblck: <p>The <code>default_collate</code> function of your <code>DataLoader</code> is creating the <code>list</code>:</p>
<pre><code class="lang-python">dataset = torch.utils.data.TensorDataset(data_geant_eval)
print(type(dataset[0]))
# &lt;class 'tuple'&gt;

torch.utils.data._utils.collate.default_collate([dataset[0], dataset[1]])
# [tensor([1., 1.])]
</code></pre>
<p>while the dataset does return the <code>tuple</code>.<br>
You can find the code <a href="https://github.com/pytorch/pytorch/blob/a17a7ccc92144452e0fe51e02f21f1f1ba88118a/torch/utils/data/_utils/collate.py">here</a>.</p>

lichuyang03: <p>I also met this problem. And I also use Tensordataset to feed in dataloader. So you have the solution now? Data in dataloader is still a list, but I wanna a tuple, how can I do that?</p>
RuntimeError: Unexpected error from cudaGetDeviceCount(); Error 802: system not yet initialized

nfarisha: <p>Hi all,</p>
<p>I’m trying to set up the environment for CUDA before importing <code>torch</code>, but I’m running into the following error:</p>
<p><code>RuntimeError: Unexpected error from cudaGetDeviceCount(). Did you run some CUDA functions before calling NumCudaDevices() that might have already set an error? Error 802: system not yet initialized.</code></p>
<p>Here’s the code I’m using:</p>
<pre><code class="lang-auto">import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"

import torch
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("Device count:", torch.cuda.device_count())
    print("Torch version:", torch.__version__)
    print("Current device:", torch.cuda.current_device())
    print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
</code></pre>
<p>Am I setting these environment variables correctly? Is there something I’m missing that could cause <code>cudaGetDeviceCount()</code> to fail with Error 802?</p>
<p>Any insights would be appreciated!</p>
<p>Thanks in advance</p>

ptrblck: <p>What does <code>deviceQuery</code> return? Assuming the same error, check if your system uses NVSwitch and if FabricManager was properly installed.</p>
Use `torch.Tensor` as typehint

pascalm: <p>Hello,</p>
<p>I would like to use <code>torch.Tensor</code> as a type hint. E.g.</p>
<pre data-code-wrap="py"><code class="lang-py">import torch

data: torch.Tensor = torch.tensor(2)
</code></pre>
<p>but if I do that I get:</p>
<p><code>foo.py:1: error: Cannot find implementation or library stub for module named "torch"</code></p>
<p>I saw there’s <a href="https://github.com/patrick-kidger/jaxtyping" class="inline-onebox" rel="noopener nofollow ugc">GitHub - patrick-kidger/jaxtyping: Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/</a> that can be used for even more in depth typing but I think for now a very basic type like the above would be sufficient for me.</p>
<p>So before I dig into jaxtyping: Is it possible to just use <code>torch.Tensor</code> as a type?</p>
Adding parentheses to boolean tensor creates small error

the-ml-bulldog: <p>Hi, I am currently trying to replicate PyTorch’s <code>nn.CrossEntropyLoss</code> module from scratch. I have the following code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

class CELoss(nn.Module):
  def __init__(self, ignore_index=-100, reduction='mean'):
    super().__init__()
    self.ignore_index = ignore_index
    self.reduction = reduction

  def forward(self, y_pred, y):
    y_pred = torch.log_softmax(y_pred, dim=1)
    unmask = y != self.ignore_index
    y = y[unmask]
    y_pred = y_pred.transpose(0, 1)
    y_pred = y_pred[:, unmask]
    indices_list = (y, torch.arange(y.shape[0]))
    if self.reduction == 'mean':
      return torch.mean(-y_pred[indices_list])
    elif self.reduction == 'sum':
      return torch.sum(-y_pred[indices_list])
    else:
      return -y_pred[indices_list]

C = 4
N = 100
size = (32, 64)
ignore_index = -100

input = torch.randn(size=(N, C, *size))
target = torch.randint(0, C, size=(N, *size), dtype=torch.long)
mask = torch.randn_like(target.to(torch.float32)) &lt; -1

target[mask] = ignore_index

reduction = "sum"
my_celoss = CELoss(reduction=reduction)
nn_celoss = nn.CrossEntropyLoss(reduction=reduction)

loss1 = my_celoss(input, target)
loss2 = nn_celoss(input, target)
#print(loss1)
#print(loss2)
print(torch.abs(loss1 - loss2)) # Comparison
</code></pre>
<p>The output is 0.0 when I set <code>unmask = y != self.ignore_index</code> in line 12, but when I set <code>unmask = (y != self.ignore_index)</code> (with parentheses), the output is around 0.03.</p>
<p>Why would adding parentheses here affect the answer?</p>

albanD: <p>Hey!</p>
<p>Running this locally gives 0. difference in both cases for me.<br>
Can you try and print the result of the boolean op to see what the difference is?</p>
<p>Note that some of these ops might be non-deterministic so small difference between runs is expected.</p>
"turing_fp16_s1688gemm_fp16_128x128_ldg8_relu_f2f_tn"

JB13: <p>What exactly this kernel does ? And why this kernel got launched during QKV projection? Are there some internal working which actually deactivate this ReLu parameter of this kernel for matrix multiplication?</p>
<p>FYI: It is launched by CPU kernel : aten::addmm</p>

ptrblck: <p>This is a <code>matmul</code> kernel for <code>float16</code> data used in mixed-precision workloads, most likely called by cuBLAS.</p>

JB13: <p>Okay but why there is this ‘relu’ ? there is no need of relu activation function in linear projection right??</p>

ptrblck: <p>The kernel has the ability to use a <code>relu</code> in the epilogue but the name itself does not mean it’s necessarily used. You can check its output and would know for sure if it’s used or not.</p>

JB13: <p>Thanks, Can you tell me how exactly I can check the output of this kernel. Because in the trace json file, I can see the output dimension. It is doing only the projection for Q,K and V matrices by multiplying the weight matrix with the Input matrix.</p>

ptrblck: <p>Isolate the <code>matmul</code> call which uses this kernel e.g. via profiling your code with Nsight Systems with stacktraces, then print the <code>matmul</code> output.</p>
Higher Peak Memory on First Run with torch.compile and Inductor Backend

nd21: <p>I am using <code>torch.compile</code> with the Inductor backend to optimize a LLaMA-2-7B model. I’ve noticed a consistent pattern: when I run the compiled model end-to-end, the peak GPU memory usage is always higher on the first run compared to subsequent runs (second, third, etc.).</p>
<p>This happens even though the model and inputs are identical each time. The memory delta between the first and second run is always is approximately (seq_len × seq_len × 2) bytes but I have not completely verified this for multiple data points.</p>
<p>For example, with seq_len = 8192, the memory drops by approximately 128 MiB between the first and second run. The behavior is deterministic and reproducible.</p>
<p>Unfortunately, I’m unable to share the codebase. But I am working with the HuggingFace implementation of LLaMA-2-7B, compiled using torch.compile and inductor backend.</p>
<p>I’m trying to understand what might be causing this delta — is it related to temporary allocations during graph capture, causal mask creation, or workspace buffers that are reused/cached in later runs? Any pointers on how to systematically debug this would be greatly appreciated.</p>
<p>Thanks!</p>
Using tensors as indices

DDong: <p>Why is it that the following returns a tensor with shape [10,4]?<br>
<code><br>
test=torch.ones(10,4)<br>
print(test[torch.arange(10)][torch.arange(4)].shape)<br>
</code><br>
Intuitively it should return the entire original tensor, test.</p>

ptrblck: <p>You are indexing in dim0 twice. This code would yield the same result:</p>
<pre><code class="lang-python">test=torch.randn(10,4)
a = test[torch.arange(10)]
b = a[torch.arange(4)]
print((test[:4] == b).all())
</code></pre>

DDong: <p>How can I index across dim1 then?</p>

ptrblck: <p>If you want to use an index tensor (e.g. [0, 1]) for all elements in dim0, this would work:</p>
<pre><code class="lang-python">test=torch.randn(10,4)
idx = torch.tensor([0, 1])
test[:, idx]
</code></pre>

DDong: <p>But what I want to happen is to select certain indices for each dimension, that is if I were to do<br>
<code><br>
test=torch.ones(10,4)<br>
print(test[torch.arange(9),torch.arange(3)])<br>
</code><br>
then it would print a 2-D tensor of ones of size [9,3]; that is, the first 9 rows, and for each of those rows, the first 3 elements. What would be the equivalent functionality? The above code returns<br>
<code><br>
IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [9], [3]<br>
</code></p>

cloudydory: <aside class="quote no-group" data-username="DDong" data-post="5" data-topic="53421">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/ec9cab/48.png" class="avatar"> DDong:</div>
<blockquote>
<p>print(test[torch.arange(9),torch.arange(3)])</p>
</blockquote>
</aside>
<p>I still meet this error in PyTorch 2.6.0</p>