Improve Neural Tangent Kernels tutorial (#2476)

- Allow it to work on CPU
- Disable TensorFloat on Ampere+ GPU to meet accuracy expectations
- Add TensorFloat to dictionary of known words

Author: malfet

en-wordlist.txt

@@ -154,6 +154,7 @@ TPU
 TensorBoard
 TensorBoards
 TensorDict
+TensorFloat
 TextVQA
 Tokenization
 TorchDynamo

intermediate_source/neural_tangent_kernels.py

@@ -24,7 +24,7 @@
 import torch
 import torch.nn as nn
 from torch.func import functional_call, vmap, vjp, jvp, jacrev
-device = 'cuda'
+device = 'cuda' if torch.cuda.device_count() > 0 else 'cpu'
 
 class CNN(nn.Module):
     def __init__(self):
@@ -224,8 +224,11 @@ def get_ntk_slice(vec):
     if compute == 'diagonal':
         return torch.einsum('NMKK->NMK', result)
 
-result_from_jacobian_contraction = empirical_ntk_jacobian_contraction(fnet_single, params, x_test, x_train)
-result_from_ntk_vps = empirical_ntk_ntk_vps(fnet_single, params, x_test, x_train)
+# Disable TensorFloat-32 for convolutions on Ampere+ GPUs to sacrifice performance in favor of accuracy
+with torch.backends.cudnn.flags(allow_tf32=False):
+    result_from_jacobian_contraction = empirical_ntk_jacobian_contraction(fnet_single, params, x_test, x_train)
+    result_from_ntk_vps = empirical_ntk_ntk_vps(fnet_single, params, x_test, x_train)
+
 assert torch.allclose(result_from_jacobian_contraction, result_from_ntk_vps, atol=1e-5)
 
 ######################################################################