{ "source_file": "0930", "source_platform": "Medium", "author": "Chaim Rand", "title_listed": "Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch", "keywords": "", "domain": "数据处理", "content_type": "深度技术解析", "is_paid_flag": "否", "url": "https://towardsdatascience.com/overcoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch/", "publish_date": "2025-12-03", "corpus_stage": "", "notes": "", "id": "48fc99ba02a3", "collected_at": "2026-10-02T11:13:01", "title": "Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch | Towards Data Science", "body_markdown": "![Revisit consent button](../images/body-df29beacd7496220cc94.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\n# Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch\nPyTorch Model Performance Analysis and Optimization — Part 11\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nDecember 3, 20259 min read\n![](../images/000.jpg)\n[Photo](https://unsplash.com/photos/red-and-black-stars-in-the-sky-zz_3tCcrk7o) by [National Cancer Institute](https://unsplash.com/@nci) on Unsplash\nThis is the eleventh part of a [series of posts](https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-10c3c5822869/) on the topic of analyzing and optimizing PyTorch models. Throughout the series, we have advocated for using the [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) in AI model development and demonstrated the potential impact of performance optimization on the speed and cost of running AI/ML workloads. One common phenomenon we have seen is how seemingly innocent code can hamper runtime performance. In this post, we explore some of the penalties associated with the naive use of variable-shaped tensors — tensors whose shape is dependent on preceding computations and/or inputs. While not applicable to all situations, there are times when the use of variable-shaped tensors can be avoided — although this may come at the expense of additional compute and/or memory. We will demonstrate the tradeoffs of these alternatives on a toy implementation of data sampling in PyTorch.\n## Three Downsides of Variable Shaped Tensors\nWe motivate the discussion by presenting three disadvantages to the use of variable-shaped tensors:\nLearn this step by step with the interactive [Machine Learning](https://roadmap.sh/machine-learning) roadmap.\n### Host-Device Sync Events\nIn an ideal scenario, the CPU and GPU are able to run in parallel in an asynchronous manner, with the CPU continuously feeding the GPU with input samples, allocating required GPU memory, and loading GPU compute kernels, and the GPU executing the loaded kernels on the provided inputs using the allocated memory. The presence of dynamic-shaped tensors throws a wrench into this parallelism. In order to allocate the appropriate amount memory, the CPU must wait for the GPU to report the tensor’s shape, and then the GPU must wait for the CPU to allocate the memory and proceed with the kernel loading. The overhead of this sync event can cause a drop in the GPU utilization and slow runtime performance.\nWe saw an example of this in [part three](https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-part-3-1c5876d78fe2) of this series when we studied a naive implementation of the common [cross-entropy loss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) that included calls to [torch.nonzero](https://pytorch.org/docs/stable/generated/torch.nonzero.html) and [torch.unique](https://pytorch.org/docs/stable/generated/torch.unique.html). Both APIs return tensors with shapes that are dynamic and dependent on the contents of the input. When these functions are run on the GPU, a host-device synchronization event occurs. In the case of the [cross-entropy loss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html), we discovered the inefficiency through the use of [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) and were able to easily overcome it with an alternative implementation that avoided the use of variable-shaped tensors and demonstrated much better runtime performance.\n### Graph Compilation\nIn a [recent post](https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/) we explored the performance benefits of applying _just-in-time_ (JIT) compilation using the [torch.compile](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) operator. One of our observations was that graph compilation provided much better results when the graph was static. The presence of dynamic shapes in the graph limits the extent of the optimization via compilation: In some cases, it fails completely; in others it results in lower performance gains. The same implications also apply to other forms of graph compilation, such as [XLA](https://openxla.org/xla), [ONNX](https://onnx.ai/), [OpenVINO](https://github.com/openvinotoolkit/openvino), and [TensorRT](https://developer.nvidia.com/tensorrt).\n### Data Batching\nAnother optimization we have encountered in several of our posts (e.g., [here](https://chaimrand.medium.com/optimizing-pytorch-model-inference-on-cpu-ccd3aa5884ad)) is sample-batching. Batching improves performance in two primary ways:\n 1. **Reducing overhead of kernel loading** : Rather than loading the GPU kernels required for the computation pipeline once per input sample, the CPU can load the kernels once per batch.\n 2. **Maximizing parallelization across compute units** : GPUs are highly parallel compute engines. The more we are able to parallelize computation, the more we can saturate the GPU and increase its utilization. By batching we can potentially increase the degree of parallelization by a factor of the batch size.\n\n\nDespite their downsides, the use of variable-shaped tensors is often unavoidable. But sometimes we can modify our model implementation to circumvent them. Sometimes these changes will be straightforward (as in the cross-entropy loss example). Other times they may require some creativity in coming up with a different sequence of fixed-shape PyTorch APIs that provide the same numerical result. Often, this effort can deliver meaningful rewards in runtime and costs.\nIn the next sections, we will study the use of variable-shaped tensors in the context of the data sampling operation. We will start with a trivial implementation and analyze its performance. We will then propose a GPU-friendly alternative that avoids the use of variable-shaped tensors.\nTo compare our implementations, we will use an [Amazon EC2 g6e.xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) with an [NVIDIA L40S](https://www.nvidia.com/en-eu/data-center/l40s/) running an [AWS Deep Learning AMI (DLAMI)](https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html) with [PyTorch (2.8)](https://pytorch.org/get-started/locally/). The code we will share is intended for demonstration purposes. Please do not rely on it for accuracy or optimality. Please do not interpret our mention of any framework, library, or platform and an endorsement of its use.\n## Sampling in AI Model Workloads\nIn the context of this post, sampling refers to the selection of a subset of items from a large set of candidates for the purposes of computational efficiency, balancing of datatypes, or regularization. Sampling is common in many AI/ML models, such as detection, ranking, and contrastive learning systems.\nWe define a simple variation of the sampling problem: Given a list of _N_ tensors each with a binary label, we are asked to return a subset of _K_ tensors containing both positive and negative examples, in random order. If the input list contains enough samples of each label (_K/2_), the returned subset should be evenly split. If it is lacking samples of one type, these should be filled with random samples of the second type.\nThe code block below contains a PyTorch implementation of our sampling function. The implementation is inspired by the popular [Detectron2](https://github.com/facebookresearch/detectron2) library (e.g., see [here](https://github.com/facebookresearch/detectron2/blob/v0.6/detectron2/modeling/sampling.py#L9) and [here](https://github.com/facebookresearch/detectron2/blob/v0.6/detectron2/modeling/roi_heads/roi_heads.py#L181)). For the experiments in this post, we will fix the sampling ratio to _1:10._\npython\n\n```\nimport torch\n\nINPUT_SAMPLES = 10000\nSUB_SAMPLE = INPUT_SAMPLES // 10\nFEATURE_DIM = 16\n\ndef sample_data(input_array, labels):\n device = labels.device\n positive = torch.nonzero(labels == 1, as_tuple=True)[0]\n negative = torch.nonzero(labels == 0, as_tuple=True)[0]\n num_pos = min(positive.numel(), SUB_SAMPLE//2)\n num_neg = min(negative.numel(), SUB_SAMPLE//2)\n if num_neg < SUB_SAMPLE//2:\n num_pos = SUB_SAMPLE - num_neg\n elif num_pos < SUB_SAMPLE//2:\n num_neg = SUB_SAMPLE - num_pos\n\n # randomly select positive and negative examples\n perm1 = torch.randperm(positive.numel(), device=device)[:num_pos]\n perm2 = torch.randperm(negative.numel(), device=device)[:num_neg]\n\n pos_idxs = positive[perm1]\n neg_idxs = negative[perm2]\n\n sampled_idxs = torch.cat([pos_idxs, neg_idxs], dim=0)\n rand_perm = torch.randperm(SUB_SAMPLE, device=labels.device)\n sampled_idxs = sampled_idxs[rand_perm]\n return input_array[sampled_idxs], labels[sampled_idxs]\n```\n\n## Performance Analysis With PyTorch Profiler\nEven when not immediately obvious, the use of dynamic shapes is easily identifiable in the PyTorch Profiler Trace view. We use the following function to enable PyTorch Profiler:\npython\n\n```\ndef profile(fn, input, labels):\n \n def export_trace(p):\n p.export_chrome_trace(f\"{fn.__name__}.json\")\n \n with torch.profiler.profile(\n activities=[torch.profiler.ProfilerActivity.CPU,\n torch.profiler.ProfilerActivity.CUDA],\n with_stack=True,\n schedule=torch.profiler.schedule(wait=0, warmup=10, active=5),\n on_trace_ready=export_trace\n ) as prof:\n for _ in range(20):\n fn(input, labels)\n torch.cuda.synchronize() # explicit sync for trace readability\n prof.step()\n\n# create random input\ninput_samples = torch.randn((INPUT_SAMPLES, FEATURE_DIM), device='cuda')\nlabels = torch.randint(0, 2, (INPUT_SAMPLES,), \n device='cuda', dtype=torch.int64)\n\n# run with profiler\nprofile(sample_data, input_samples, labels)\n```\n\nThe image below was captured for the value of ten million input samples. It clearly shows the presence of sync events coming from the torch.nonzero call, as well as the corresponding drops in GPU utilization:\n![](../images/001.webp)\nProfiler Trace of Sampler (by Author)\nThe use of torch.nonzero in our implementation is not ideal, but can it be avoided?\n## A GPU-Friendly Data Sampler\nWe propose an alternative implementation of our sampling function that replaces the dynamic torch.nonzero function with a creative combination of the static [torch.count_nonzero](https://docs.pytorch.org/docs/stable/generated/torch.count_nonzero.html), [torch.topk](https://docs.pytorch.org/docs/stable/generated/torch.topk.html), and other APIs:\npython\n\n```\ndef opt_sample_data(input, labels):\n pos_mask = labels == 1\n neg_mask = labels == 0\n num_pos_idxs = torch.count_nonzero(pos_mask, dim=-1)\n num_neg_idxs = torch.count_nonzero(neg_mask, dim=-1)\n half_samples = labels.new_full((), SUB_SAMPLE // 2)\n num_pos = torch.minimum(num_pos_idxs, half_samples)\n num_neg = torch.minimum(num_neg_idxs, half_samples)\n num_pos = torch.where(\n num_neg < SUB_SAMPLE // 2,\n SUB_SAMPLE - num_neg,\n num_pos\n )\n num_neg = SUB_SAMPLE - num_pos\n\n # create random ordering on pos and neg entries\n rand = torch.rand_like(labels, dtype=torch.float32)\n pos_rand = torch.where(pos_mask, rand, -1)\n neg_rand = torch.where(neg_mask, rand, -1)\n\n # select top pos entries and invalidate others\n # since CPU doesn't know num_pos, we assume maximum to avoid sync\n top_pos_rand, top_pos_idx = torch.topk(pos_rand, k=SUB_SAMPLE)\n arange = torch.arange(SUB_SAMPLE, device=labels.device)\n if num_pos.numel() > 1:\n # unsqueeze to support batched input\n arange = arange.unsqueeze(0)\n num_pos = num_pos.unsqueeze(-1)\n num_neg = num_neg.unsqueeze(-1)\n top_pos_rand = torch.where(arange >= num_pos, -1, top_pos_rand)\n\n # repeat for neg entries\n top_neg_rand, top_neg_idx = torch.topk(neg_rand, k=SUB_SAMPLE)\n top_neg_rand = torch.where(arange >= num_neg, -1, top_neg_rand)\n\n # combine and mix together positive and negative idxs\n cat_rand = torch.cat([top_pos_rand, top_neg_rand], dim=-1)\n cat_idx = torch.cat([top_pos_idx, top_neg_idx], dim=-1)\n topk_rand_idx = torch.topk(cat_rand, k=SUB_SAMPLE)[1]\n sampled_idxs = torch.gather(cat_idx, dim=-1, index=topk_rand_idx)\n sampled_input = torch.gather(input, dim=-2, \n index=sampled_idxs.unsqueeze(-1))\n sampled_labels = torch.gather(labels, dim=-1, index=sampled_idxs)\n return sampled_input, sampled_labels\n```\n\nClearly, this function requires more memory and more operations than our first implementation. The question is: Do the performance benefits of a static, synchronization-free implementation outweigh the extra cost in memory and compute?\nTo assess the tradeoffs between the two implementations, we introduce the following benchmarking utility:\npython\n\n```\ndef benchmark(fn, input, labels):\n # warm-up\n for _ in range(20):\n _ = fn(input, labels)\n\n iters = 100\n start = torch.cuda.Event(enable_timing=True)\n end = torch.cuda.Event(enable_timing=True)\n torch.cuda.synchronize()\n start.record()\n for _ in range(iters):\n _ = fn(input, labels)\n end.record()\n torch.cuda.synchronize()\n avg_time = start.elapsed_time(end) / iters\n \n print(f\"{fn.__name__} average step time: {(avg_time):.4f} ms\")\n\nbenchmark(sample_data, input_samples, labels)\nbenchmark(opt_sample_data, input_samples, labels)\n```\n\nThe following table compares the average runtime of each of the implementations for a variety of input sample sizes:\n![](../images/body-3cc35a457c68bd4b7655.webp)\nComparative Step Time Performance — Lower is Better (by Author)\nFor most of the input sample sizes, the overhead of the host-device sync event is either comparable or lower than the additional compute of the static implementation. Disappointingly, we only see a major benefit from the sync-free alternative when the input sample size reaches ten million. Sample sizes that large are uncommon in AI/ML settings. But it’s not our tendency to give up so easily. As noted above, the static implementation enables other optimizations like graph compilation and input batching.\n### Graph Compilation\nContrary to the original function — which fails to compile — our static implementation is fully compatible with torch.compile:\npython\n\n```\nbenchmark(torch.compile(opt_sample_data), input_samples, labels)\n```\n\nThe following table includes the runtimes of our compiled function:\n![](../images/body-70a21ed6458a0af66717.webp)\nComparative Step Time Performance — Lower is Better (by Author)\nThe results are significantly better — providing a 70–75 percent boost over the original sampler implementation in the 1–10 thousand range. But we still have one more optimization up our sleeve.\n### Maximizing Performance with Batched Input\nBecause the original implementation contains variable-shaped operations, it cannot handle batched input directly. To process a batch, we have no choice but to apply it to each input individually, in a Python loop:\npython\n\n```\nBATCH_SIZE = 32\n\ndef batched_sample_data(inputs, labels):\n sampled_inputs = []\n sampled_labels = []\n for i in range(inputs.size(0)):\n inp, lab = sample_data(inputs[i], labels[i])\n sampled_inputs.append(inp)\n sampled_labels.append(lab)\n return torch.stack(sampled_inputs), torch.stack(sampled_labels)\n```\n\nIn contrast, our optimized function supports batched inputs as is — no changes necessary.\npython\n\n```\ninput_batch = torch.randn((BATCH_SIZE, INPUT_SAMPLES, FEATURE_DIM),\n device='cuda')\nlabels = torch.randint(0, 2, (BATCH_SIZE, INPUT_SAMPLES),\n device='cuda', dtype=torch.int64)\n\nbenchmark(batched_sample_data, input_batch, labels)\nbenchmark(opt_sample_data, input_batch, labels)\nbenchmark(torch.compile(opt_sample_data), input_batch, labels)\n```\n\nThe table below compares the step times of our sampling functions on a batch size of 32:\n![](../images/body-e2cd3db2ea7722b59a34.webp)\nStep Time Performance on Batched Input — Lower is Better (by Author)\nNow the results are definitive: By using a static implementation of the data sampler, we are able to boost performance by 2X–52X(!!) the variable-shaped option, depending on the input sample size.\nNote that although our experiments were run on a GPU device, the model compilation and input batching optimizations also apply to a CPU environment. Thus, avoiding variable shapes could have implications on AI/ML model performance on CPU, as well.\n## Summary\nThe optimization process we demonstrated in this post generalizes beyond the specific case of data sampling:\n * **Discovery via Performance Profiling:** Using the [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) we were able to identify drops in GPU utilization and discover their source: the presence of variable-shaped tensors resulting from the torch.nonzero operation.\n * **An Alternate Implementation:** Our profiling findings allowed us to develop an alternative implementation that accomplished the same goal while avoiding the use of variable-shaped tensors. However, this step came at the cost of additional compute and memory overhead. As seen in our initial benchmarks, the sync-free alternative demonstrated worse performance on common input sizes.\n * **Unlocking Further Potential for Optimization:** The true breakthrough came because the static-shaped implementation was compilation-friendly and supported batching. These optimizations provided performance gains that dwarfed the initial overhead, leading to a 2x to 52x speedup over the original implementation.\n\n\nNaturally, not all stories will end as happily as ours. In many cases, we may come across PyTorch code that performs poorly on the GPU but does not have an alternative implementation, or it may have one that requires significantly more compute resources. However, given the potential for meaningful gains in performance and reductions in cost, the process of identifying runtime inefficiencies and exploring alternative implementations is an essential part of AI/ML development.\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](../images/body-97f5483d57a9b8741146.jpeg)\nWritten By\n## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Data Sampling](https://towardsdatascience.com/tag/data-sampling/)[Machine Learning](https://towardsdatascience.com/tag/machine-learning/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fovercoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch%2F&title=Overcoming%20the%20Hidden%20Performance%20Traps%20of%20Variable-Shaped%20Tensors%3A%20Efficient%20Data%20Sampling%20in%20PyTorch)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fovercoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch%2F&title=Overcoming%20the%20Hidden%20Performance%20Traps%20of%20Variable-Shaped%20Tensors%3A%20Efficient%20Data%20Sampling%20in%20PyTorch)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fovercoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch%2F&text=Overcoming%20the%20Hidden%20Performance%20Traps%20of%20Variable-Shaped%20Tensors%3A%20Efficient%20Data%20Sampling%20in%20PyTorch)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](../images/002.jpg)](https://towardsdatascience.com/from-resume-to-cover-letter-using-ai-and-llm-with-python-and-streamlit/)\n## [From Resume to Cover Letter Using AI and LLM, with Python and Streamlit](https://towardsdatascience.com/from-resume-to-cover-letter-using-ai-and-llm-with-python-and-streamlit/)\n[Large Language Models](https://towardsdatascience.com/category/artificial-intelligence/large-language-models/)\nHere’s how less than 100 lines of code can build you an LLM app that converts your resume into your job-specific cover letter\n[Piero Paialunga](https://towardsdatascience.com/author/piero-paialunga/)\nFebruary 5, 202513 min read\n * [![](../images/003.jpeg)](https://towardsdatascience.com/your-vision-language-model-might-be-a-bag-of-words-30b1beaef7f8/)\n## [Your Vision-Language Model Might Be a Bag of Words](https://towardsdatascience.com/your-vision-language-model-might-be-a-bag-of-words-30b1beaef7f8/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nWe explore the limits of what vision-language models get about language in our Oral Paper at ICLR 2023\n[Federico Bianchi](https://towardsdatascience.com/author/fede-bianchi/)\nMarch 21, 20239 min read\n * [![](../images/004.jpeg)](https://towardsdatascience.com/hyperparameter-optimization-intro-and-implementation-of-grid-search-random-search-and-bayesian-b2f16c00578a/)\n## [Hyperparameter Optimization - Intro and Implementation of Grid Search, Random Search and Bayesian...](https://towardsdatascience.com/hyperparameter-optimization-intro-and-implementation-of-grid-search-random-search-and-bayesian-b2f16c00578a/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nMost common hyperparameter optimization methodologies to boost machine learning outcomes.\n[Farzad Nobar](https://towardsdatascience.com/author/fmnobar/)\nMarch 13, 202310 min read\n * [![](../images/005.jpeg)](https://towardsdatascience.com/analysis-of-us-elections-from-1976-to-2010-with-pandas-6654cf69849d/)\n## [Analysis of US Elections from 1976 to 2010 with Pandas](https://towardsdatascience.com/analysis-of-us-elections-from-1976-to-2010-with-pandas-6654cf69849d/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nPractice Pandas while providing valuable insights into past elections.\n[Soner Yıldırım](https://towardsdatascience.com/author/sonery/)\nOctober 18, 20205 min read\n * [![](../images/006.jpeg)](https://towardsdatascience.com/8-must-have-git-commands-for-data-scientists-ee6564e4631d/)\n## [8 Must-Have Git Commands for Data Scientists](https://towardsdatascience.com/8-must-have-git-commands-for-data-scientists-ee6564e4631d/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nThis is how you contribute to a project\n[Soner Yıldırım](https://towardsdatascience.com/author/sonery/)\nJuly 18, 20214 min read\n * [![](../images/007.jpg)](https://towardsdatascience.com/why-your-ai-search-evaluation-is-probably-wrong-and-how-to-fix-it/)\n## [Why Your AI Search Evaluation Is Probably Wrong (And How to Fix It)](https://towardsdatascience.com/why-your-ai-search-evaluation-is-probably-wrong-and-how-to-fix-it/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nA five-step framework for building rigorous, reproducible AI search benchmarks — before you make six-figure infrastructure decisions\n[Zairah Mustahsan](https://towardsdatascience.com/author/zairah.mustahsan/)\nMarch 9, 20266 min read\n * [![](../images/008.jpeg)](https://towardsdatascience.com/fit-vs-transform-in-scikit-libraries-for-machine-learning-3c70e6300ded/)\n## [Fit vs. Transform in SciKit libraries for Machine Learning](https://towardsdatascience.com/fit-vs-transform-in-scikit-libraries-for-machine-learning-3c70e6300ded/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nWe have seen methods such as fit(), transform(), and fit_transform() in a lot of SciKit's libraries. And almost all tutorials, including...\n[Sunny Srinidhi](https://towardsdatascience.com/author/contactsunny/)\nDecember 3, 20193 min read\n * [![](../images/009.jpg)](https://towardsdatascience.com/5-tips-that-will-help-you-become-a-data-science-manager-bd4b6981c024/)\n## [5 Tips That Will Help You Become a Data Science Manager](https://towardsdatascience.com/5-tips-that-will-help-you-become-a-data-science-manager-bd4b6981c024/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nSo you want to be promoted? So does everyone else. How can you stand out from your peers?\n[Tessa Xie](https://towardsdatascience.com/author/tessaxie/)\nMay 24, 20226 min read\n * [![](../images/010.jpeg)](https://towardsdatascience.com/intergenerational-mobility-in-the-us-a-data-science-problem-1-5-5ab0d6569edd/)\n## [Intergenerational Mobility in the US - A Data Science Problem (1/5)](https://towardsdatascience.com/intergenerational-mobility-in-the-us-a-data-science-problem-1-5-5ab0d6569edd/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nThis series was written together with Ibukun Aribilola and Valdrin Jonuzi as part of the Data Science for Social Good tutorial.\n[Ang Li-Lian](https://towardsdatascience.com/author/anglilian/)\nApril 22, 202211 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "metadata": { "title": "Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch | Towards Data Science", "description": "PyTorch Model Performance Analysis and Optimization — Part 11", "keywords": null, "author": "Chaim Rand", "og:locale": "en_US", "og:type": "article", "og:title": "Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch", "og:description": "PyTorch Model Performance Analysis and Optimization — Part 11", "og:url": "https://towardsdatascience.com/overcoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch/", "og:site_name": "Towards Data Science", "og:image": "https://assets.insightmediagroup.io/media/wp-content/uploads/2025/12/national-cancer-institute-zz_3tCcrk7o-unsplash.jpg", "og:image:width": "1920", "og:image:height": "1367", "og:image:type": "image/jpeg", "twitter:card": "summary_large_image", "twitter:site": "@TDataScience", "twitter:creator": "@TDataScience", "twitter:label1": "Written by", "twitter:data1": "Chaim Rand", "twitter:label2": "Est. reading time", "twitter:data2": "9 min read", "article:published_time": "2025-12-03T12:00:06.000Z", "article:modified_time": "2025-12-03T12:00:06.000Z", "article:author": "https://towardsdatascience.com/author/chaimrand/" }, "paywall": { "listed_paid": false, "detected": false, "marker": "", "public_only": false, "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙" }, "crawl": { "success": true, "status_code": 202, "error": "" }, "code_blocks": [ { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_000.txt", "chars": 1001 }, { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_001.txt", "chars": 864 }, { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_002.txt", "chars": 1933 }, { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_003.txt", "chars": 582 }, { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_004.txt", "chars": 65 }, { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_005.txt", "chars": 333 }, { "language": "txt", "file": "attachments/48fc99ba02a3/code/block_006.txt", "chars": 388 } ], "images": [ { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:06:18Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:07:10Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:07:10Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:07:10Z" } ], "audio": [], "videos": [ { "url": "https://www.youtube.com/c/TowardsDataScience", "kind": "channel_reference", "status": "reference_only" } ], "links": [ { "type": "internal", "href": "https://towardsdatascience.com/" }, { "type": "internal", "href": "https://towardsdatascience.com/latest/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/editors-pick/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/the-variable/" }, { "type": "internal", "href": "https://towardsdatascience.com/shipai/" }, { "type": "internal", "href": "https://towardsdatascience.com/submissions/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/deep-learning/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/chaimrand/" }, { "type": "internal", "href": "https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-10c3c5822869/" }, { "type": "internal", "href": "https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-part-3-1c5876d78fe2" }, { "type": "internal", "href": "https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/artificial-intelligence/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/data-sampling/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/machine-learning/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/pytorch/" }, { "type": "internal", "href": "https://towardsdatascience.com/from-resume-to-cover-letter-using-ai-and-llm-with-python-and-streamlit/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/large-language-models/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/piero-paialunga/" }, { "type": "internal", "href": "https://towardsdatascience.com/your-vision-language-model-might-be-a-bag-of-words-30b1beaef7f8/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/" }, { "type": "external", "href": "https://write.insightmediagroup.io" }, { "type": "external", "href": "https://www.linkedin.com/company/towards-data-science/" }, { "type": "external", "href": "https://x.com/TDataScience" }, { "type": "external", "href": "https://unsplash.com/photos/red-and-black-stars-in-the-sky-zz_3tCcrk7o" }, { "type": "external", "href": "https://unsplash.com/@nci" }, { "type": "external", "href": "https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html" }, { "type": "external", "href": "https://roadmap.sh/machine-learning" }, { "type": "external", "href": "https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html" }, { "type": "external", "href": "https://pytorch.org/docs/stable/generated/torch.nonzero.html" }, { "type": "external", "href": "https://pytorch.org/docs/stable/generated/torch.unique.html" }, { "type": "external", "href": "https://docs.pytorch.org/docs/stable/generated/torch.compile.html" }, { "type": "external", "href": "https://openxla.org/xla" }, { "type": "external", "href": "https://onnx.ai/" }, { "type": "external", "href": "https://github.com/openvinotoolkit/openvino" }, { "type": "external", "href": "https://developer.nvidia.com/tensorrt" }, { "type": "external", "href": "https://chaimrand.medium.com/optimizing-pytorch-model-inference-on-cpu-ccd3aa5884ad" }, { "type": "external", "href": "https://aws.amazon.com/ec2/instance-types/g6e/" }, { "type": "external", "href": "https://www.nvidia.com/en-eu/data-center/l40s/" }, { "type": "external", "href": "https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html" }, { "type": "external", "href": "https://pytorch.org/get-started/locally/" } ], "repair_version": 2, "repaired_at": "2026-10-20T17:32:37Z", "source_url": "https://towardsdatascience.com/overcoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch/", "delivery_status": "unverified", "discovery_status": "http_403", "missing_reason": "源页面未重新核验:http_403；已按保留正文补采图片", "repair_error": "", "actual_video_files": [], "expected_saved_body_images": 4, "image_discovery": "saved_markdown", "source_body_markdown": "![Revisit consent button](https://cdn-cookieyes.com/assets/images/revisit.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](https://cdn-cookieyes.com/assets/images/close.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\n# Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch\nPyTorch Model Performance Analysis and Optimization — Part 11\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nDecember 3, 20259 min read\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/12/national-cancer-institute-zz_3tCcrk7o-unsplash.jpg)\n[Photo](https://unsplash.com/photos/red-and-black-stars-in-the-sky-zz_3tCcrk7o) by [National Cancer Institute](https://unsplash.com/@nci) on Unsplash\nThis is the eleventh part of a [series of posts](https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-10c3c5822869/) on the topic of analyzing and optimizing PyTorch models. Throughout the series, we have advocated for using the [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) in AI model development and demonstrated the potential impact of performance optimization on the speed and cost of running AI/ML workloads. One common phenomenon we have seen is how seemingly innocent code can hamper runtime performance. In this post, we explore some of the penalties associated with the naive use of variable-shaped tensors — tensors whose shape is dependent on preceding computations and/or inputs. While not applicable to all situations, there are times when the use of variable-shaped tensors can be avoided — although this may come at the expense of additional compute and/or memory. We will demonstrate the tradeoffs of these alternatives on a toy implementation of data sampling in PyTorch.\n## Three Downsides of Variable Shaped Tensors\nWe motivate the discussion by presenting three disadvantages to the use of variable-shaped tensors:\nLearn this step by step with the interactive [Machine Learning](https://roadmap.sh/machine-learning) roadmap.\n### Host-Device Sync Events\nIn an ideal scenario, the CPU and GPU are able to run in parallel in an asynchronous manner, with the CPU continuously feeding the GPU with input samples, allocating required GPU memory, and loading GPU compute kernels, and the GPU executing the loaded kernels on the provided inputs using the allocated memory. The presence of dynamic-shaped tensors throws a wrench into this parallelism. In order to allocate the appropriate amount memory, the CPU must wait for the GPU to report the tensor’s shape, and then the GPU must wait for the CPU to allocate the memory and proceed with the kernel loading. The overhead of this sync event can cause a drop in the GPU utilization and slow runtime performance.\nWe saw an example of this in [part three](https://towardsdatascience.com/pytorch-model-performance-analysis-and-optimization-part-3-1c5876d78fe2) of this series when we studied a naive implementation of the common [cross-entropy loss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) that included calls to [torch.nonzero](https://pytorch.org/docs/stable/generated/torch.nonzero.html) and [torch.unique](https://pytorch.org/docs/stable/generated/torch.unique.html). Both APIs return tensors with shapes that are dynamic and dependent on the contents of the input. When these functions are run on the GPU, a host-device synchronization event occurs. In the case of the [cross-entropy loss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html), we discovered the inefficiency through the use of [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) and were able to easily overcome it with an alternative implementation that avoided the use of variable-shaped tensors and demonstrated much better runtime performance.\n### Graph Compilation\nIn a [recent post](https://towardsdatascience.com/maximizing-ai-ml-model-performance-with-pytorch-compilation/) we explored the performance benefits of applying _just-in-time_ (JIT) compilation using the [torch.compile](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) operator. One of our observations was that graph compilation provided much better results when the graph was static. The presence of dynamic shapes in the graph limits the extent of the optimization via compilation: In some cases, it fails completely; in others it results in lower performance gains. The same implications also apply to other forms of graph compilation, such as [XLA](https://openxla.org/xla), [ONNX](https://onnx.ai/), [OpenVINO](https://github.com/openvinotoolkit/openvino), and [TensorRT](https://developer.nvidia.com/tensorrt).\n### Data Batching\nAnother optimization we have encountered in several of our posts (e.g., [here](https://chaimrand.medium.com/optimizing-pytorch-model-inference-on-cpu-ccd3aa5884ad)) is sample-batching. Batching improves performance in two primary ways:\n 1. **Reducing overhead of kernel loading** : Rather than loading the GPU kernels required for the computation pipeline once per input sample, the CPU can load the kernels once per batch.\n 2. **Maximizing parallelization across compute units** : GPUs are highly parallel compute engines. The more we are able to parallelize computation, the more we can saturate the GPU and increase its utilization. By batching we can potentially increase the degree of parallelization by a factor of the batch size.\n\n\nDespite their downsides, the use of variable-shaped tensors is often unavoidable. But sometimes we can modify our model implementation to circumvent them. Sometimes these changes will be straightforward (as in the cross-entropy loss example). Other times they may require some creativity in coming up with a different sequence of fixed-shape PyTorch APIs that provide the same numerical result. Often, this effort can deliver meaningful rewards in runtime and costs.\nIn the next sections, we will study the use of variable-shaped tensors in the context of the data sampling operation. We will start with a trivial implementation and analyze its performance. We will then propose a GPU-friendly alternative that avoids the use of variable-shaped tensors.\nTo compare our implementations, we will use an [Amazon EC2 g6e.xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) with an [NVIDIA L40S](https://www.nvidia.com/en-eu/data-center/l40s/) running an [AWS Deep Learning AMI (DLAMI)](https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html) with [PyTorch (2.8)](https://pytorch.org/get-started/locally/). The code we will share is intended for demonstration purposes. Please do not rely on it for accuracy or optimality. Please do not interpret our mention of any framework, library, or platform and an endorsement of its use.\n## Sampling in AI Model Workloads\nIn the context of this post, sampling refers to the selection of a subset of items from a large set of candidates for the purposes of computational efficiency, balancing of datatypes, or regularization. Sampling is common in many AI/ML models, such as detection, ranking, and contrastive learning systems.\nWe define a simple variation of the sampling problem: Given a list of _N_ tensors each with a binary label, we are asked to return a subset of _K_ tensors containing both positive and negative examples, in random order. If the input list contains enough samples of each label (_K/2_), the returned subset should be evenly split. If it is lacking samples of one type, these should be filled with random samples of the second type.\nThe code block below contains a PyTorch implementation of our sampling function. The implementation is inspired by the popular [Detectron2](https://github.com/facebookresearch/detectron2) library (e.g., see [here](https://github.com/facebookresearch/detectron2/blob/v0.6/detectron2/modeling/sampling.py#L9) and [here](https://github.com/facebookresearch/detectron2/blob/v0.6/detectron2/modeling/roi_heads/roi_heads.py#L181)). For the experiments in this post, we will fix the sampling ratio to _1:10._\npython\n\n```\nimport torch\n\nINPUT_SAMPLES = 10000\nSUB_SAMPLE = INPUT_SAMPLES // 10\nFEATURE_DIM = 16\n\ndef sample_data(input_array, labels):\n device = labels.device\n positive = torch.nonzero(labels == 1, as_tuple=True)[0]\n negative = torch.nonzero(labels == 0, as_tuple=True)[0]\n num_pos = min(positive.numel(), SUB_SAMPLE//2)\n num_neg = min(negative.numel(), SUB_SAMPLE//2)\n if num_neg < SUB_SAMPLE//2:\n num_pos = SUB_SAMPLE - num_neg\n elif num_pos < SUB_SAMPLE//2:\n num_neg = SUB_SAMPLE - num_pos\n\n # randomly select positive and negative examples\n perm1 = torch.randperm(positive.numel(), device=device)[:num_pos]\n perm2 = torch.randperm(negative.numel(), device=device)[:num_neg]\n\n pos_idxs = positive[perm1]\n neg_idxs = negative[perm2]\n\n sampled_idxs = torch.cat([pos_idxs, neg_idxs], dim=0)\n rand_perm = torch.randperm(SUB_SAMPLE, device=labels.device)\n sampled_idxs = sampled_idxs[rand_perm]\n return input_array[sampled_idxs], labels[sampled_idxs]\n```\n\n## Performance Analysis With PyTorch Profiler\nEven when not immediately obvious, the use of dynamic shapes is easily identifiable in the PyTorch Profiler Trace view. We use the following function to enable PyTorch Profiler:\npython\n\n```\ndef profile(fn, input, labels):\n \n def export_trace(p):\n p.export_chrome_trace(f\"{fn.__name__}.json\")\n \n with torch.profiler.profile(\n activities=[torch.profiler.ProfilerActivity.CPU,\n torch.profiler.ProfilerActivity.CUDA],\n with_stack=True,\n schedule=torch.profiler.schedule(wait=0, warmup=10, active=5),\n on_trace_ready=export_trace\n ) as prof:\n for _ in range(20):\n fn(input, labels)\n torch.cuda.synchronize() # explicit sync for trace readability\n prof.step()\n\n# create random input\ninput_samples = torch.randn((INPUT_SAMPLES, FEATURE_DIM), device='cuda')\nlabels = torch.randint(0, 2, (INPUT_SAMPLES,), \n device='cuda', dtype=torch.int64)\n\n# run with profiler\nprofile(sample_data, input_samples, labels)\n```\n\nThe image below was captured for the value of ten million input samples. It clearly shows the presence of sync events coming from the torch.nonzero call, as well as the corresponding drops in GPU utilization:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/11/2-1-1024x198.webp)\nProfiler Trace of Sampler (by Author)\nThe use of torch.nonzero in our implementation is not ideal, but can it be avoided?\n## A GPU-Friendly Data Sampler\nWe propose an alternative implementation of our sampling function that replaces the dynamic torch.nonzero function with a creative combination of the static [torch.count_nonzero](https://docs.pytorch.org/docs/stable/generated/torch.count_nonzero.html), [torch.topk](https://docs.pytorch.org/docs/stable/generated/torch.topk.html), and other APIs:\npython\n\n```\ndef opt_sample_data(input, labels):\n pos_mask = labels == 1\n neg_mask = labels == 0\n num_pos_idxs = torch.count_nonzero(pos_mask, dim=-1)\n num_neg_idxs = torch.count_nonzero(neg_mask, dim=-1)\n half_samples = labels.new_full((), SUB_SAMPLE // 2)\n num_pos = torch.minimum(num_pos_idxs, half_samples)\n num_neg = torch.minimum(num_neg_idxs, half_samples)\n num_pos = torch.where(\n num_neg < SUB_SAMPLE // 2,\n SUB_SAMPLE - num_neg,\n num_pos\n )\n num_neg = SUB_SAMPLE - num_pos\n\n # create random ordering on pos and neg entries\n rand = torch.rand_like(labels, dtype=torch.float32)\n pos_rand = torch.where(pos_mask, rand, -1)\n neg_rand = torch.where(neg_mask, rand, -1)\n\n # select top pos entries and invalidate others\n # since CPU doesn't know num_pos, we assume maximum to avoid sync\n top_pos_rand, top_pos_idx = torch.topk(pos_rand, k=SUB_SAMPLE)\n arange = torch.arange(SUB_SAMPLE, device=labels.device)\n if num_pos.numel() > 1:\n # unsqueeze to support batched input\n arange = arange.unsqueeze(0)\n num_pos = num_pos.unsqueeze(-1)\n num_neg = num_neg.unsqueeze(-1)\n top_pos_rand = torch.where(arange >= num_pos, -1, top_pos_rand)\n\n # repeat for neg entries\n top_neg_rand, top_neg_idx = torch.topk(neg_rand, k=SUB_SAMPLE)\n top_neg_rand = torch.where(arange >= num_neg, -1, top_neg_rand)\n\n # combine and mix together positive and negative idxs\n cat_rand = torch.cat([top_pos_rand, top_neg_rand], dim=-1)\n cat_idx = torch.cat([top_pos_idx, top_neg_idx], dim=-1)\n topk_rand_idx = torch.topk(cat_rand, k=SUB_SAMPLE)[1]\n sampled_idxs = torch.gather(cat_idx, dim=-1, index=topk_rand_idx)\n sampled_input = torch.gather(input, dim=-2, \n index=sampled_idxs.unsqueeze(-1))\n sampled_labels = torch.gather(labels, dim=-1, index=sampled_idxs)\n return sampled_input, sampled_labels\n```\n\nClearly, this function requires more memory and more operations than our first implementation. The question is: Do the performance benefits of a static, synchronization-free implementation outweigh the extra cost in memory and compute?\nTo assess the tradeoffs between the two implementations, we introduce the following benchmarking utility:\npython\n\n```\ndef benchmark(fn, input, labels):\n # warm-up\n for _ in range(20):\n _ = fn(input, labels)\n\n iters = 100\n start = torch.cuda.Event(enable_timing=True)\n end = torch.cuda.Event(enable_timing=True)\n torch.cuda.synchronize()\n start.record()\n for _ in range(iters):\n _ = fn(input, labels)\n end.record()\n torch.cuda.synchronize()\n avg_time = start.elapsed_time(end) / iters\n \n print(f\"{fn.__name__} average step time: {(avg_time):.4f} ms\")\n\nbenchmark(sample_data, input_samples, labels)\nbenchmark(opt_sample_data, input_samples, labels)\n```\n\nThe following table compares the average runtime of each of the implementations for a variety of input sample sizes:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/11/3-1.webp)\nComparative Step Time Performance — Lower is Better (by Author)\nFor most of the input sample sizes, the overhead of the host-device sync event is either comparable or lower than the additional compute of the static implementation. Disappointingly, we only see a major benefit from the sync-free alternative when the input sample size reaches ten million. Sample sizes that large are uncommon in AI/ML settings. But it’s not our tendency to give up so easily. As noted above, the static implementation enables other optimizations like graph compilation and input batching.\n### Graph Compilation\nContrary to the original function — which fails to compile — our static implementation is fully compatible with torch.compile:\npython\n\n```\nbenchmark(torch.compile(opt_sample_data), input_samples, labels)\n```\n\nThe following table includes the runtimes of our compiled function:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/11/4-1.webp)\nComparative Step Time Performance — Lower is Better (by Author)\nThe results are significantly better — providing a 70–75 percent boost over the original sampler implementation in the 1–10 thousand range. But we still have one more optimization up our sleeve.\n### Maximizing Performance with Batched Input\nBecause the original implementation contains variable-shaped operations, it cannot handle batched input directly. To process a batch, we have no choice but to apply it to each input individually, in a Python loop:\npython\n\n```\nBATCH_SIZE = 32\n\ndef batched_sample_data(inputs, labels):\n sampled_inputs = []\n sampled_labels = []\n for i in range(inputs.size(0)):\n inp, lab = sample_data(inputs[i], labels[i])\n sampled_inputs.append(inp)\n sampled_labels.append(lab)\n return torch.stack(sampled_inputs), torch.stack(sampled_labels)\n```\n\nIn contrast, our optimized function supports batched inputs as is — no changes necessary.\npython\n\n```\ninput_batch = torch.randn((BATCH_SIZE, INPUT_SAMPLES, FEATURE_DIM),\n device='cuda')\nlabels = torch.randint(0, 2, (BATCH_SIZE, INPUT_SAMPLES),\n device='cuda', dtype=torch.int64)\n\nbenchmark(batched_sample_data, input_batch, labels)\nbenchmark(opt_sample_data, input_batch, labels)\nbenchmark(torch.compile(opt_sample_data), input_batch, labels)\n```\n\nThe table below compares the step times of our sampling functions on a batch size of 32:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/11/5.webp)\nStep Time Performance on Batched Input — Lower is Better (by Author)\nNow the results are definitive: By using a static implementation of the data sampler, we are able to boost performance by 2X–52X(!!) the variable-shaped option, depending on the input sample size.\nNote that although our experiments were run on a GPU device, the model compilation and input batching optimizations also apply to a CPU environment. Thus, avoiding variable shapes could have implications on AI/ML model performance on CPU, as well.\n## Summary\nThe optimization process we demonstrated in this post generalizes beyond the specific case of data sampling:\n * **Discovery via Performance Profiling:** Using the [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) we were able to identify drops in GPU utilization and discover their source: the presence of variable-shaped tensors resulting from the torch.nonzero operation.\n * **An Alternate Implementation:** Our profiling findings allowed us to develop an alternative implementation that accomplished the same goal while avoiding the use of variable-shaped tensors. However, this step came at the cost of additional compute and memory overhead. As seen in our initial benchmarks, the sync-free alternative demonstrated worse performance on common input sizes.\n * **Unlocking Further Potential for Optimization:** The true breakthrough came because the static-shaped implementation was compilation-friendly and supported batching. These optimizations provided performance gains that dwarfed the initial overhead, leading to a 2x to 52x speedup over the original implementation.\n\n\nNaturally, not all stories will end as happily as ours. In many cases, we may come across PyTorch code that performs poorly on the GPU but does not have an alternative implementation, or it may have one that requires significantly more compute resources. However, given the potential for meaningful gains in performance and reductions in cost, the process of identifying runtime inefficiencies and exploring alternative implementations is an essential part of AI/ML development.\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](https://assets.insightmediagroup.io/media/1787657094227_7ymyo0.jpeg)\nWritten By\n## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Data Sampling](https://towardsdatascience.com/tag/data-sampling/)[Machine Learning](https://towardsdatascience.com/tag/machine-learning/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fovercoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch%2F&title=Overcoming%20the%20Hidden%20Performance%20Traps%20of%20Variable-Shaped%20Tensors%3A%20Efficient%20Data%20Sampling%20in%20PyTorch)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fovercoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch%2F&title=Overcoming%20the%20Hidden%20Performance%20Traps%20of%20Variable-Shaped%20Tensors%3A%20Efficient%20Data%20Sampling%20in%20PyTorch)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fovercoming-the-hidden-performance-traps-of-variable-shaped-tensors-efficient-data-sampling-in-pytorch%2F&text=Overcoming%20the%20Hidden%20Performance%20Traps%20of%20Variable-Shaped%20Tensors%3A%20Efficient%20Data%20Sampling%20in%20PyTorch)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/02/1_KX5gpc1KEypqkCOll-HDtw-scaled.jpg)](https://towardsdatascience.com/from-resume-to-cover-letter-using-ai-and-llm-with-python-and-streamlit/)\n## [From Resume to Cover Letter Using AI and LLM, with Python and Streamlit](https://towardsdatascience.com/from-resume-to-cover-letter-using-ai-and-llm-with-python-and-streamlit/)\n[Large Language Models](https://towardsdatascience.com/category/artificial-intelligence/large-language-models/)\nHere’s how less than 100 lines of code can build you an LLM app that converts your resume into your job-specific cover letter\n[Piero Paialunga](https://towardsdatascience.com/author/piero-paialunga/)\nFebruary 5, 202513 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/03/1xGMde9uuUeHNIRShrkBRSA-scaled.jpeg)](https://towardsdatascience.com/your-vision-language-model-might-be-a-bag-of-words-30b1beaef7f8/)\n## [Your Vision-Language Model Might Be a Bag of Words](https://towardsdatascience.com/your-vision-language-model-might-be-a-bag-of-words-30b1beaef7f8/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nWe explore the limits of what vision-language models get about language in our Oral Paper at ICLR 2023\n[Federico Bianchi](https://towardsdatascience.com/author/fede-bianchi/)\nMarch 21, 20239 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/03/1NUiYEKQ6Ad8mzSUJKhUphw-scaled.jpeg)](https://towardsdatascience.com/hyperparameter-optimization-intro-and-implementation-of-grid-search-random-search-and-bayesian-b2f16c00578a/)\n## [Hyperparameter Optimization - Intro and Implementation of Grid Search, Random Search and Bayesian...](https://towardsdatascience.com/hyperparameter-optimization-intro-and-implementation-of-grid-search-random-search-and-bayesian-b2f16c00578a/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nMost common hyperparameter optimization methodologies to boost machine learning outcomes.\n[Farzad Nobar](https://towardsdatascience.com/author/fmnobar/)\nMarch 13, 202310 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2020/10/19meZjVfvGg2qDJq3ZD4bzA-scaled.jpeg)](https://towardsdatascience.com/analysis-of-us-elections-from-1976-to-2010-with-pandas-6654cf69849d/)\n## [Analysis of US Elections from 1976 to 2010 with Pandas](https://towardsdatascience.com/analysis-of-us-elections-from-1976-to-2010-with-pandas-6654cf69849d/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nPractice Pandas while providing valuable insights into past elections.\n[Soner Yıldırım](https://towardsdatascience.com/author/sonery/)\nOctober 18, 20205 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/07/1_tp5YBaWh8aSKhEHHmkcoA-scaled.jpeg)](https://towardsdatascience.com/8-must-have-git-commands-for-data-scientists-ee6564e4631d/)\n## [8 Must-Have Git Commands for Data Scientists](https://towardsdatascience.com/8-must-have-git-commands-for-data-scientists-ee6564e4631d/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nThis is how you contribute to a project\n[Soner Yıldırım](https://towardsdatascience.com/author/sonery/)\nJuly 18, 20214 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/03/a-chosen-soul-D_ivYIn4jWw-unsplash.jpg)](https://towardsdatascience.com/why-your-ai-search-evaluation-is-probably-wrong-and-how-to-fix-it/)\n## [Why Your AI Search Evaluation Is Probably Wrong (And How to Fix It)](https://towardsdatascience.com/why-your-ai-search-evaluation-is-probably-wrong-and-how-to-fix-it/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nA five-step framework for building rigorous, reproducible AI search benchmarks — before you make six-figure infrastructure decisions\n[Zairah Mustahsan](https://towardsdatascience.com/author/zairah.mustahsan/)\nMarch 9, 20266 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2019/12/1EfCv3Z60nG1wsc3_Cf_Zsg.jpeg)](https://towardsdatascience.com/fit-vs-transform-in-scikit-libraries-for-machine-learning-3c70e6300ded/)\n## [Fit vs. Transform in SciKit libraries for Machine Learning](https://towardsdatascience.com/fit-vs-transform-in-scikit-libraries-for-machine-learning-3c70e6300ded/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nWe have seen methods such as fit(), transform(), and fit_transform() in a lot of SciKit's libraries. And almost all tutorials, including...\n[Sunny Srinidhi](https://towardsdatascience.com/author/contactsunny/)\nDecember 3, 20193 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/05/0eEkocPod0idXkFTo-scaled.jpg)](https://towardsdatascience.com/5-tips-that-will-help-you-become-a-data-science-manager-bd4b6981c024/)\n## [5 Tips That Will Help You Become a Data Science Manager](https://towardsdatascience.com/5-tips-that-will-help-you-become-a-data-science-manager-bd4b6981c024/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nSo you want to be promoted? So does everyone else. How can you stand out from your peers?\n[Tessa Xie](https://towardsdatascience.com/author/tessaxie/)\nMay 24, 20226 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/04/1_XbsNgvJWdjTU-GloDS0tQ-scaled.jpeg)](https://towardsdatascience.com/intergenerational-mobility-in-the-us-a-data-science-problem-1-5-5ab0d6569edd/)\n## [Intergenerational Mobility in the US - A Data Science Problem (1/5)](https://towardsdatascience.com/intergenerational-mobility-in-the-us-a-data-science-problem-1-5-5ab0d6569edd/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nThis series was written together with Ibukun Aribilola and Valdrin Jonuzi as part of the Data Science for Social Good tutorial.\n[Ang Li-Lian](https://towardsdatascience.com/author/anglilian/)\nApril 22, 202211 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "body_base_path": "attachments/48fc99ba02a3/pages/", "body_markdown_path": "attachments/48fc99ba02a3/pages/archived-body.md", "local_html": "attachments/48fc99ba02a3/pages/archived-body.html", "body_source": "prior_capture" }