pkgs/development/cuda-modules/packages/tensorrt-samples/testers/sample_dynamic_reshape.nix


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

{
  atLeast,
  lib,
  mkTester,
  sample-data,
  ...
}:
{
  default = mkTester "sample_dynamic_reshape" [
    "sample_dynamic_reshape"
    "--datadir=${sample-data.outPath + "/mnist"}"
  ];

  # TODO(@connorbaker): Neither the sample data nor the sample sources provide train-images-idx3-ubyte, so we can't run the int8 test.
  # int8 = mkTester "sample_dynamic_reshape-int8" [
  #   "sample_dynamic_reshape"
  #   "--datadir=${sample-data.outPath + "/mnist"}"
  #   "--int8"
  # ];

  fp16 = mkTester "sample_dynamic_reshape-fp16" [
    "sample_dynamic_reshape"
    "--datadir=${sample-data.outPath + "/mnist"}"
    "--fp16"
  ];
}
// lib.optionalAttrs (atLeast "10") {
  bf16 = mkTester "sample_dynamic_reshape-bf16" [
    "sample_dynamic_reshape"
    "--datadir=${sample-data.outPath + "/mnist"}"
    "--bf16"
  ];
}