diff --git a/components/swap/SwapFormTokenList.tsx b/components/swap/SwapFormTokenList.tsx
index e1e5ecda..efcb788a 100644
--- a/components/swap/SwapFormTokenList.tsx
+++ b/components/swap/SwapFormTokenList.tsx
@@ -125,13 +125,15 @@ const SwapFormTokenList = ({
       group &&
       mangoAccount &&
       outputBank &&
+      inputBank &&
       type === 'input'
     ) {
       const filteredSortedTokens = mangoTokens
         .map((token) => {
           const max = getTokenInMax(
             mangoAccount,
-            token.address,
+            inputBank.mint,
+            outputBank.mint,
             group,
             useMargin
           )
diff --git a/components/swap/useTokenMax.tsx b/components/swap/useTokenMax.tsx
index 1b7859b2..bbfdf4d1 100644
--- a/components/swap/useTokenMax.tsx
+++ b/components/swap/useTokenMax.tsx
@@ -5,6 +5,7 @@ import mangoStore from '@store/mangoStore'
 import { floorToDecimal } from '../../utils/numbers'
 import useMangoAccount from '../../hooks/useMangoAccount'
 import useMangoGroup from 'hooks/useMangoGroup'
+import { PublicKey } from '@solana/web3.js'
 
 export const getMaxWithdrawForBank = (
   group: Group,
@@ -26,12 +27,13 @@ export const getMaxWithdrawForBank = (
 
 export const getTokenInMax = (
   mangoAccount: MangoAccount,
-  inputTokenAddress: string,
+  inputMint: PublicKey,
+  outputMint: PublicKey,
   group: Group,
   useMargin: boolean
 ) => {
-  const outputBank = mangoStore.getState().swap.outputBank
-  const inputBank = group.banksMapByMint.get(inputTokenAddress)?.[0]
+  const inputBank = group.getFirstBankByMint(inputMint)
+  const outputBank = group.getFirstBankByMint(outputMint)
 
   if (!group || !inputBank || !mangoAccount || !outputBank) {
     return {
@@ -72,16 +74,6 @@ export const getTokenInMax = (
       )
     : Decimal.min(maxAmountWithoutMargin, inputBankVaultBalance)
 
-  console.log(
-    'maxUiAmountWithBorrow',
-    mangoAccount.getMaxSourceUiForTokenSwap(
-      group,
-      inputBank.mint,
-      outputBank.mint,
-      1
-    )
-  )
-
   const maxAmountWithBorrow = Decimal.min(
     maxUiAmountWithBorrow,
     inputBankVaultBalance
@@ -98,12 +90,14 @@ export const useTokenMax = (useMargin = true) => {
   const { mangoAccount } = useMangoAccount()
   const { group } = useMangoGroup()
   const inputBank = mangoStore((s) => s.swap.inputBank)
+  const outputBank = mangoStore((s) => s.swap.outputBank)
 
   const tokenInMax = useMemo(() => {
-    if (mangoAccount && group && inputBank) {
+    if (mangoAccount && group && inputBank && outputBank) {
       return getTokenInMax(
         mangoAccount,
-        inputBank?.mint.toString(),
+        inputBank.mint,
+        outputBank.mint,
         group,
         useMargin
       )
@@ -114,7 +108,7 @@ export const useTokenMax = (useMargin = true) => {
       amountWithBorrow: new Decimal(0),
       decimals: 6,
     }
-  }, [mangoAccount, group, useMargin, inputBank])
+  }, [mangoAccount, group, useMargin, inputBank, outputBank])
 
   return tokenInMax
 }